feat: add insert batch and streaming rate controls

.env.example

@@ -2,7 +2,6 @@ LOG_LEVEL=INFO
 LOG_FILE="logs/vectordb_bench.log"
 # TIMEZONE=
 
-# NUM_PER_BATCH=
 # DEFAULT_DATASET_URL=
 
 DATASET_LOCAL_DIR="/tmp/vectordb_bench/dataset"

README.md

@@ -287,7 +287,7 @@ OpenSearch Serverless (AOSS) is a serverless deployment option for Amazon OpenSe
 **Example: Run performance test on OpenSearch Serverless**
 
 ```shell
-NUM_PER_BATCH=100 vectordbbench awsopensearch --db-label aoss \
+vectordbbench awsopensearch --db-label aoss --insert-batch-size 100 \
   --serverless --aws-region us-east-1 \
   --host <collection-id>.aoss.us-east-1.on.aws --port 443 \
   --case-type Performance768D1M \
@@ -303,13 +303,13 @@ OpenSearch Serverless-specific options:
 |--------|-------------|
 | `--serverless` | Enable OpenSearch Serverless mode (uses AWS SigV4 auth) |
 | `--aws-region` | AWS region for the AOSS collection (default: `us-east-1`) |
-| `NUM_PER_BATCH` | Number of vectors per Serverless bulk request (default: `100`) |
+| `--insert-batch-size` | Number of vectors per Serverless bulk request (default: `100`) |
 
 > **Notes:**
 > - `--user` and `--password` are not needed for Serverless mode
 > - `--engine` is accepted but ignored internally (AOSS manages the engine)
 > - `--force-merge-enabled`, `--refresh-interval`, `--flush-threshold-size`, and `--cb-threshold` are ignored for Serverless
-> - Keep `NUM_PER_BATCH` small enough for the Serverless bulk API request limits
+> - Keep `--insert-batch-size` small enough for the Serverless bulk API request limits
 
 ### Run Elastic Cloud from command line
 
@@ -478,7 +478,7 @@ pip install 'vectordb-bench[hologres]' 'psycopg[binary]' pgvector
 Execute tests for the index types: HGraph.
 
 ```shell
-NUM_PER_BATCH=10000 vectordbbench hologreshgraph --host Hologres_Endpoint --port 80 \
+vectordbbench hologreshgraph --host Hologres_Endpoint --port 80 --insert-batch-size 10000 \
 --user ACCESS_ID --password ACCESS_KEY --database DATABASE_NAME \
 --m 64 --ef-construction 400 --case-type Performance768D10M \
 --index-type HGraph --ef-search 400 --k 10 --num-concurrency 1,60,70,75,80,90,95,100,105,110,115,120,125,130 \
@@ -532,12 +532,12 @@ To list the options for zvec, execute vectordbbench zvec --help
 Doris supports ann index with type hnsw from version 4.0.x
 
 ```shell
-NUM_PER_BATCH=1000000 vectordbbench doris --http-port=8030 --port=9030 --db-name=vector_test --case-type=Performance768D1M --stream-load-rows-per-batch=500000
+vectordbbench doris --http-port=8030 --port=9030 --db-name=vector_test --case-type=Performance768D1M --insert-batch-size=1000000 --stream-load-rows-per-batch=500000
 ```
 
 Using flag `--session-var`, if you want to test doris with some customized session variables. For example:
 ```shell
-NUM_PER_BATCH=1000000 vectordbbench doris --http-port=8030 --port=9030 --db-name=vector_test --case-type=Performance768D1M --stream-load-rows-per-batch=500000 --session-var enable_profile=True
+vectordbbench doris --http-port=8030 --port=9030 --db-name=vector_test --case-type=Performance768D1M --insert-batch-size=1000000 --stream-load-rows-per-batch=500000 --session-var enable_profile=True
 ```
 
 Mote options:
@@ -558,8 +558,8 @@ Mote options:
 --session-var TEXT              Session variable key=value applied to each
                                 SQL session (repeatable)
 --stream-load-rows-per-batch INTEGER
-                                Rows per single stream load request; default
-                                uses NUM_PER_BATCH
+                                Rows per Doris stream-load request; when
+                                omitted, the Doris client default is used
 --no-index                      Create table without ANN index
 ```
 

docs/release/2026-05-cloud-leaderboard.md

@@ -65,7 +65,7 @@ vectordbbench zillizautoindex \
   --uri "$ZILLIZ_URI" \
   --token "$ZILLIZ_TOKEN" \
   --collection-name cloud_insert_laion100m_bs10k \
-  --cloud-insert-batch-size 10000 \
+  --insert-batch-size 10000 \
   --load-concurrency 16 \
   --skip-search-serial \
   --skip-search-concurrent \

tests/test_aws_opensearch.py

@@ -2,21 +2,19 @@
 
 import pytest
 
-from vectordb_bench import config
 from vectordb_bench.backend.clients.aws_opensearch.aws_opensearch import AWSOpenSearch
 
 
-def test_serverless_insert_uses_configured_batch_size(monkeypatch) -> None:
+def test_serverless_insert_uses_configured_batch_size() -> None:
     bulk_requests = []
 
-    def bulk(*, body):
+    def bulk(*, body: list) -> None:
         bulk_requests.append(body)
 
-    monkeypatch.setattr(config, "NUM_PER_BATCH", 2)
-
     db = object.__new__(AWSOpenSearch)
     db.client = SimpleNamespace(bulk=bulk)
     db._is_serverless = True
+    db._insert_batch_size = 2
     db.index_name = "test-index"
     db.vector_col_name = "embedding"
     db.with_scalar_labels = False
@@ -33,13 +31,12 @@ def bulk(*, body):
 
 
 @pytest.mark.parametrize("batch_size", [0, -1])
-def test_serverless_insert_rejects_non_positive_batch_size(monkeypatch, batch_size: int) -> None:
-    monkeypatch.setattr(config, "NUM_PER_BATCH", batch_size)
-
+def test_serverless_insert_rejects_non_positive_batch_size(batch_size: int) -> None:
     db = object.__new__(AWSOpenSearch)
     db._is_serverless = True
+    db._insert_batch_size = batch_size
 
-    with pytest.raises(ValueError, match="NUM_PER_BATCH must be greater than 0"):
+    with pytest.raises(ValueError, match="insert_batch_size must be greater than 0"):
         db._insert_with_single_client(
             embeddings=[[0.1]],
             metadata=[1],

tests/test_case_runner_reuse.py

@@ -1,5 +1,6 @@
 from pydantic import SecretStr
 
+from vectordb_bench import config
 from vectordb_bench.backend.clients import DB
 from vectordb_bench.backend.clients.api import EmptyDBCaseConfig, MetricType
 from vectordb_bench.backend.clients.doris.config import DorisCaseConfig, DorisConfig
@@ -12,6 +13,8 @@
 from vectordb_bench.metric import Metric
 from vectordb_bench.models import CaseConfig, CaseType, TaskConfig, TaskStage, TestResult
 
+DEFAULT_INSERT_BATCH_SIZE = config.DEFAULT_INSERT_BATCH_SIZE
+
 
 def make_runner(
     *,
@@ -21,6 +24,7 @@ def make_runner(
     db_config=None,
     db_case_config=None,
     stages: list[TaskStage] | None = None,
+    insert_batch_size: int = DEFAULT_INSERT_BATCH_SIZE,
 ) -> CaseRunner:
     if db_config is None:
         if db == DB.TurboPuffer:
@@ -45,6 +49,7 @@ def make_runner(
         db_case_config=db_case_config,
         case_config=CaseConfig(case_id=case_id, custom_case=custom_case or {}),
         stages=stages or [TaskStage.DROP_OLD, TaskStage.LOAD, TaskStage.SEARCH_SERIAL],
+        insert_batch_size=insert_batch_size,
     )
     return CaseRunner(
         run_id="run-id",
@@ -111,6 +116,13 @@ def test_reuse_key_preserves_safe_payload_reuse():
     assert hash(ids_only) == hash(vector)
 
 
+def test_reuse_key_distinguishes_insert_batch_size():
+    assert_not_reusable(
+        make_runner(insert_batch_size=100),
+        make_runner(insert_batch_size=200),
+    )
+
+
 def test_reuse_key_distinguishes_physical_db_targets():
     assert_not_reusable(
         make_runner(db_config=TurboPufferConfig(api_key="key", region="aws-us-east-1", namespace="namespace_a")),

tests/test_cloud_insert_case.py

@@ -69,13 +69,13 @@ def iter_batches(self, batch_size):
 
 
 def test_cloud_insert_case_defaults_to_laion_100m():
-    case = CloudInsertCase(batch_size=1000)
+    case = CloudInsertCase()
 
     assert case.case_id == CaseType.CloudInsertCase
     assert case.label == CaseLabel.CloudInsert
     assert case.dataset.data.name == "LAION"
     assert case.dataset.data.size == 100_000_000
-    assert case.batch_size == 1000
+    assert not hasattr(case, "batch_size")
     assert case.duration is None
     assert case.readiness_timeout is None
 
@@ -84,14 +84,13 @@ def test_case_config_builds_cloud_insert_case_from_custom_case():
     case = CaseConfig(
         case_id=CaseType.CloudInsertCase,
         custom_case={
-            "batch_size": 5000,
             "duration": 1800,
             "dataset_with_size_type": DatasetWithSizeType.CohereMedium.value,
         },
     ).case
 
     assert isinstance(case, CloudInsertCase)
-    assert case.batch_size == 5000
+    assert not hasattr(case, "batch_size")
     assert case.duration == 1800
     assert case.dataset.data.name == "Cohere"
     assert case.dataset.data.size == 1_000_000
@@ -101,13 +100,12 @@ def test_case_config_builds_cloud_insert_case_from_laion_100m_dataset_option():
     case = CaseConfig(
         case_id=CaseType.CloudInsertCase,
         custom_case={
-            "batch_size": 10_000,
             "dataset_with_size_type": "Large LAION (768dim, 100M)",
         },
     ).case
 
     assert isinstance(case, CloudInsertCase)
-    assert case.batch_size == 10_000
+    assert not hasattr(case, "batch_size")
     assert case.dataset_with_size_type == DatasetWithSizeType.LAIONLarge
     assert case.dataset.data.name == "LAION"
     assert case.dataset.data.size == 100_000_000
@@ -122,15 +120,13 @@ def test_laion_100m_dataset_option_uses_100m_timeouts():
 def test_cli_builds_cloud_insert_custom_case_config():
     params = {
         "case_type": "CloudInsertCase",
-        "cloud_insert_batch_size": 10_000,
         "cloud_insert_duration": 1800,
         "cloud_insert_readiness_timeout": 7200,
         "cloud_insert_readiness_poll_interval": 10,
         "dataset_with_size_type": DatasetWithSizeType.CohereMedium.value,
     }
 
     assert get_custom_case_config(params) == {
-        "batch_size": 10_000,
         "duration": 1800,
         "readiness_timeout": 7200,
         "readiness_poll_interval": 10,
@@ -142,7 +138,6 @@ def test_cli_builds_cloud_insert_custom_case_config_with_laion_100m_dataset():
     cfg = get_custom_case_config(
         {
             "case_type": "CloudInsertCase",
-            "cloud_insert_batch_size": 10_000,
             "cloud_insert_duration": None,
             "cloud_insert_readiness_timeout": None,
             "cloud_insert_readiness_poll_interval": None,
@@ -151,7 +146,6 @@ def test_cli_builds_cloud_insert_custom_case_config_with_laion_100m_dataset():
     )
 
     assert cfg == {
-        "batch_size": 10_000,
         "duration": None,
         "dataset_with_size_type": DatasetWithSizeType.LAIONLarge.value,
     }
@@ -166,7 +160,6 @@ def test_cli_builds_cloud_insert_custom_case_config_with_default_dataset():
     cfg = get_custom_case_config(
         {
             "case_type": "CloudInsertCase",
-            "cloud_insert_batch_size": 10_000,
             "cloud_insert_duration": None,
             "cloud_insert_readiness_timeout": None,
             "cloud_insert_readiness_poll_interval": None,
@@ -175,7 +168,6 @@ def test_cli_builds_cloud_insert_custom_case_config_with_default_dataset():
     )
 
     assert cfg == {
-        "batch_size": 10_000,
         "duration": None,
         "dataset_with_size_type": DatasetWithSizeType.CohereMedium.value,
     }
@@ -243,11 +235,11 @@ def test_assembler_schedules_cloud_insert_case():
         case_config=CaseConfig(
             case_id=CaseType.CloudInsertCase,
             custom_case={
-                "batch_size": 1000,
                 "dataset_with_size_type": DatasetWithSizeType.CohereMedium.value,
             },
         ),
         stages=[TaskStage.DROP_OLD, TaskStage.LOAD],
+        insert_batch_size=1000,
     )
 
     runner = Assembler.assemble_all("run-id", "task-label", [task], DatasetSource.S3)
@@ -311,10 +303,11 @@ def test_cloud_insert_result_file_uses_insert_only_metrics(tmp_path: Path):
             db_case_config=EmptyDBCaseConfig(),
             case_config=CaseConfig(
                 case_id=CaseType.CloudInsertCase,
-                custom_case={"batch_size": 1000, "duration": None},
+                custom_case={"duration": None},
             ),
             stages=[TaskStage.DROP_OLD, TaskStage.LOAD],
             load_concurrency=0,
+            insert_batch_size=1000,
         ),
         metrics=Metric(
             inserted_count=100_000_000,
@@ -343,14 +336,16 @@ def test_cloud_insert_result_file_uses_insert_only_metrics(tmp_path: Path):
     }
     assert written["results"][0]["task_config"]["db_config"]["api_key"] == "**********"
     assert written["results"][0]["task_config"]["db_config"]["index_name"] == "laion100m"
+    assert written["results"][0]["task_config"]["insert_batch_size"] == 1000
     assert written["results"][0]["task_config"]["case_config"] == {
         "case_id": 600,
-        "custom_case": {"batch_size": 1000, "duration": None},
+        "custom_case": {"duration": None},
     }
 
     read_back = TestResult.read_file(result_file)
     assert read_back.results[0].task_config.case_config.case_id == CaseType.CloudInsertCase
-    assert read_back.results[0].task_config.case_config.custom_case == {"batch_size": 1000, "duration": None}
+    assert read_back.results[0].task_config.case_config.custom_case == {"duration": None}
+    assert read_back.results[0].task_config.insert_batch_size == 1000
 
     collected = ResultCollector.collect(tmp_path)
     assert len(collected) == 1
@@ -423,7 +418,7 @@ def write(self, **kwargs):
 def test_milvus_insert_readiness_uses_entity_count_and_index_progress():
     db = Milvus.__new__(Milvus)
     db.collection_name = "c"
-    db._vector_index_name = "vector_idx"
+    db._main_index_name = "vector_idx"
     db.client = type(
         "Client",
         (),
@@ -697,9 +692,9 @@ def poll_insert_readiness(self, expected_count):
 
     db = DB()
     monkeypatch.setattr("vectordb_bench.backend.task_runner.time.sleep", lambda _: None)
-    case = CloudInsertCase(batch_size=2)
+    case = CloudInsertCase()
     case.dataset = Dataset()
-    config = type("Config", (), {"load_concurrency": 1})()
+    config = type("Config", (), {"load_concurrency": 1, "insert_batch_size": 2})()
     runner = CaseRunner.construct(ca=case, db=db, config=config)
 
     metric = runner._run_cloud_insert_case()
@@ -752,9 +747,13 @@ def fail_on_sleep(_seconds):
 
     monkeypatch.setattr("vectordb_bench.backend.task_runner.ConcurrentInsertRunner", FakeConcurrentInsertRunner)
     monkeypatch.setattr("vectordb_bench.backend.task_runner.time.sleep", fail_on_sleep)
-    case = CloudInsertCase(batch_size=1, readiness_timeout=0, readiness_poll_interval=0)
+    case = CloudInsertCase(readiness_timeout=0, readiness_poll_interval=0)
     case.dataset = Dataset()
-    runner = CaseRunner.construct(ca=case, db=DB(), config=type("Config", (), {"load_concurrency": 1})())
+    runner = CaseRunner.construct(
+        ca=case,
+        db=DB(),
+        config=type("Config", (), {"load_concurrency": 1, "insert_batch_size": 1})(),
+    )
 
     with pytest.raises(TimeoutError, match="fully_searchable.*last_status.*stalled"):
         runner._run_cloud_insert_case()
@@ -798,9 +797,9 @@ def poll_insert_readiness(self, expected_count):
             return {"fully_searchable": True, "fully_indexed": True, "additional_parameters": {}}
 
     monkeypatch.setattr("vectordb_bench.backend.task_runner.ConcurrentInsertRunner", FakeConcurrentInsertRunner)
-    case = CloudInsertCase(batch_size=1000, duration=60)
+    case = CloudInsertCase(duration=60)
     case.dataset = Dataset()
-    config = type("Config", (), {"load_concurrency": 7})()
+    config = type("Config", (), {"load_concurrency": 7, "insert_batch_size": 1000})()
     runner = CaseRunner.construct(ca=case, db=DB(), config=config)
 
     metric = runner._run_cloud_insert_case()