feat: add Apache Pinot vector search client

Adds a complete Apache Pinot client for VectorDBBench.

Index types: HNSW (Lucene), IVF_FLAT, IVF_PQ, IVF_ON_DISK
Metrics: L2, IP, COSINE
Filters: NumGE, StrEqual
Optional dep: pip install "vectordb-bench[pinot]"

Parallel loading: thread_safe=True — each worker thread maintains its own
row buffer and flushes to Pinot via a fresh HTTP session. Since Pinot's
ingestFromFile is synchronous (blocks until HNSW index is built, ~6 min
per 100K×768D segment), concurrent flushes across threads reduce load time
significantly vs sequential flushing.

Benchmark results:
Small dataset (OpenAI 50K, 768D, L2):
  HNSW:        798 QPS, recall=1.000
  IVF_FLAT:    800 QPS, recall=1.000
  IVF_PQ:      795 QPS, recall=1.000
  IVF_ON_DISK: 691 QPS, recall=1.000

Large dataset (Cohere 1M, 768D, COSINE):
  HNSW m=16:    74 QPS, recall=0.982

Filter benchmark (Cohere 1M, COSINE, HNSW m=32):
  1% NumGE:  71 QPS, recall=0.977
  99% NumGE: 97 QPS, recall=0.649

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: xiangfu0

pyproject.toml

@@ -82,6 +82,7 @@ turbopuffer     = [ "turbopuffer" ]
 zvec            = [ "zvec" ]
 endee           = [ "endee==0.1.10" ]
 lindorm         = [ "opensearch-py" ]
+pinot           = [ "requests" ]
 
 [project.urls]
 Repository  = "https://github.com/zilliztech/VectorDBBench"

vectordb_bench/backend/clients/__init__.py

@@ -61,6 +61,7 @@ class DB(Enum):
     Lindorm = "Lindorm"
     VectorChord = "VectorChord"
     PolarDB = "PolarDB"
+    Pinot = "Pinot"
 
     @property
     def init_cls(self) -> type[VectorDB]:  # noqa: PLR0911, PLR0912, C901, PLR0915
@@ -257,6 +258,11 @@ def init_cls(self) -> type[VectorDB]:  # noqa: PLR0911, PLR0912, C901, PLR0915
 
             return PolarDB
 
+        if self == DB.Pinot:
+            from .pinot.pinot import Pinot
+
+            return Pinot
+
         msg = f"Unknown DB: {self.name}"
         raise ValueError(msg)
 
@@ -455,6 +461,11 @@ def config_cls(self) -> type[DBConfig]:  # noqa: PLR0911, PLR0912, C901, PLR0915
 
             return PolarDBConfig
 
+        if self == DB.Pinot:
+            from .pinot.config import PinotConfig
+
+            return PinotConfig
+
         msg = f"Unknown DB: {self.name}"
         raise ValueError(msg)
 
@@ -631,6 +642,16 @@ def case_config_cls(  # noqa: C901, PLR0911, PLR0912, PLR0915
 
             return _vectorchord_case_config.get(index_type)
 
+        if self == DB.Pinot:
+            from .pinot.config import PinotHNSWConfig, PinotIVFFlatConfig, PinotIVFOnDiskConfig, PinotIVFPQConfig
+
+            return {
+                IndexType.HNSW: PinotHNSWConfig,
+                IndexType.IVFFlat: PinotIVFFlatConfig,
+                IndexType.IVFPQ: PinotIVFPQConfig,
+                IndexType.IVFOnDisk: PinotIVFOnDiskConfig,
+            }.get(index_type, PinotHNSWConfig)
+
         # DB.Pinecone, DB.Redis
         return EmptyDBCaseConfig
 

vectordb_bench/backend/clients/api.py

@@ -25,6 +25,7 @@ class IndexType(StrEnum):
     DISKANN = "DISKANN"
     STREAMING_DISKANN = "DISKANN"
     IVFFlat = "IVF_FLAT"
+    IVFOnDisk = "IVF_ON_DISK"
     IVFPQ = "IVF_PQ"
     IVFBQ = "IVF_BQ"
     IVFSQ8 = "IVF_SQ8"

vectordb_bench/backend/clients/pinot/__init__.py

@@ -0,0 +1,250 @@
+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
+
+
+class PinotTypedDict(TypedDict):
+    controller_host: Annotated[
+        str,
+        click.option("--controller-host", type=str, default="localhost", help="Pinot Controller host"),
+    ]
+    controller_port: Annotated[
+        int,
+        click.option("--controller-port", type=int, default=9000, help="Pinot Controller port"),
+    ]
+    broker_host: Annotated[
+        str,
+        click.option("--broker-host", type=str, default="localhost", help="Pinot Broker host"),
+    ]
+    broker_port: Annotated[
+        int,
+        click.option("--broker-port", type=int, default=8099, help="Pinot Broker port"),
+    ]
+    username: Annotated[
+        str,
+        click.option("--username", type=str, default=None, help="Pinot username (optional)"),
+    ]
+    password: Annotated[
+        str,
+        click.option("--password", type=str, default=None, help="Pinot password (optional)"),
+    ]
+    ingest_batch_size: Annotated[
+        int,
+        click.option(
+            "--ingest-batch-size",
+            type=int,
+            default=100_000,
+            show_default=True,
+            help=(
+                "Rows buffered before flushing one Pinot segment (one ingestFromFile call). "
+                "Larger values mean fewer segments and better IVF training / query performance. "
+                "Reduce if memory is constrained (100K x 768-dim float32 ~= 300 MB)."
+            ),
+        ),
+    ]
+
+
+def _pinot_db_config(parameters: dict):
+    from .config import PinotConfig
+
+    return PinotConfig(
+        db_label=parameters["db_label"],
+        controller_host=parameters["controller_host"],
+        controller_port=parameters["controller_port"],
+        broker_host=parameters["broker_host"],
+        broker_port=parameters["broker_port"],
+        username=parameters.get("username"),
+        password=SecretStr(parameters["password"]) if parameters.get("password") else None,
+        ingest_batch_size=parameters["ingest_batch_size"],
+    )
+
+
+# ---------------------------------------------------------------------------
+# HNSW
+# ---------------------------------------------------------------------------
+
+
+class PinotHNSWTypedDict(CommonTypedDict, PinotTypedDict, HNSWFlavor2): ...
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(PinotHNSWTypedDict)
+def PinotHNSW(**parameters: Unpack[PinotHNSWTypedDict]):
+    from .config import PinotHNSWConfig
+
+    run(
+        db=DB.Pinot,
+        db_config=_pinot_db_config(parameters),
+        db_case_config=PinotHNSWConfig(
+            m=parameters["m"],
+            ef_construction=parameters["ef_construction"],
+            ef=parameters["ef_runtime"],
+        ),
+        **parameters,
+    )
+
+
+# ---------------------------------------------------------------------------
+# IVF_FLAT
+# ---------------------------------------------------------------------------
+
+
+class PinotIVFFlatTypedDict(CommonTypedDict, PinotTypedDict):
+    nlist: Annotated[
+        int,
+        click.option("--nlist", type=int, default=128, help="Number of Voronoi cells (IVF nlist)"),
+    ]
+    quantizer: Annotated[
+        str,
+        click.option(
+            "--quantizer",
+            type=click.Choice(["FLAT", "SQ8", "SQ4"]),
+            default="FLAT",
+            help="Quantizer type for IVF_FLAT",
+        ),
+    ]
+    nprobe: Annotated[
+        int,
+        click.option("--nprobe", type=int, default=8, help="Number of cells to probe at query time"),
+    ]
+    train_sample_size: Annotated[
+        int,
+        click.option(
+            "--train-sample-size",
+            type=int,
+            default=None,
+            help="Training sample size (defaults to max(nlist*50, 1000) if not set)",
+        ),
+    ]
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(PinotIVFFlatTypedDict)
+def PinotIVFFlat(**parameters: Unpack[PinotIVFFlatTypedDict]):
+    from .config import PinotIVFFlatConfig
+
+    run(
+        db=DB.Pinot,
+        db_config=_pinot_db_config(parameters),
+        db_case_config=PinotIVFFlatConfig(
+            nlist=parameters["nlist"],
+            quantizer=parameters["quantizer"],
+            nprobe=parameters["nprobe"],
+            train_sample_size=parameters.get("train_sample_size"),
+        ),
+        **parameters,
+    )
+
+
+# ---------------------------------------------------------------------------
+# IVF_PQ
+# ---------------------------------------------------------------------------
+
+
+class PinotIVFPQTypedDict(CommonTypedDict, PinotTypedDict):
+    nlist: Annotated[
+        int,
+        click.option("--nlist", type=int, default=128, help="Number of Voronoi cells (IVF nlist)"),
+    ]
+    pq_m: Annotated[
+        int,
+        click.option("--pq-m", type=int, default=8, help="Number of PQ sub-quantizers (must divide dimension)"),
+    ]
+    pq_nbits: Annotated[
+        int,
+        click.option(
+            "--pq-nbits",
+            type=click.Choice(["4", "6", "8"]),
+            default="8",
+            help="Bits per PQ code (4, 6, or 8)",
+        ),
+    ]
+    train_sample_size: Annotated[
+        int,
+        click.option("--train-sample-size", type=int, default=6400, help="Training sample size (must be >= nlist)"),
+    ]
+    nprobe: Annotated[
+        int,
+        click.option("--nprobe", type=int, default=8, help="Number of cells to probe at query time"),
+    ]
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(PinotIVFPQTypedDict)
+def PinotIVFPQ(**parameters: Unpack[PinotIVFPQTypedDict]):
+    from .config import PinotIVFPQConfig
+
+    run(
+        db=DB.Pinot,
+        db_config=_pinot_db_config(parameters),
+        db_case_config=PinotIVFPQConfig(
+            nlist=parameters["nlist"],
+            pq_m=parameters["pq_m"],
+            pq_nbits=int(parameters["pq_nbits"]),
+            train_sample_size=parameters["train_sample_size"],
+            nprobe=parameters["nprobe"],
+        ),
+        **parameters,
+    )
+
+
+# ---------------------------------------------------------------------------
+# IVF_ON_DISK
+# ---------------------------------------------------------------------------
+
+
+class PinotIVFOnDiskTypedDict(CommonTypedDict, PinotTypedDict):
+    nlist: Annotated[
+        int,
+        click.option("--nlist", type=int, default=128, help="Number of Voronoi cells (IVF nlist)"),
+    ]
+    quantizer: Annotated[
+        str,
+        click.option(
+            "--quantizer",
+            type=click.Choice(["FLAT", "SQ8", "SQ4"]),
+            default="FLAT",
+            help="Quantizer type for IVF_ON_DISK",
+        ),
+    ]
+    nprobe: Annotated[
+        int,
+        click.option("--nprobe", type=int, default=8, help="Number of cells to probe at query time"),
+    ]
+    train_sample_size: Annotated[
+        int,
+        click.option(
+            "--train-sample-size",
+            type=int,
+            default=None,
+            help="Training sample size (defaults to max(nlist*50, 1000) if not set)",
+        ),
+    ]
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(PinotIVFOnDiskTypedDict)
+def PinotIVFOnDisk(**parameters: Unpack[PinotIVFOnDiskTypedDict]):
+    from .config import PinotIVFOnDiskConfig
+
+    run(
+        db=DB.Pinot,
+        db_config=_pinot_db_config(parameters),
+        db_case_config=PinotIVFOnDiskConfig(
+            nlist=parameters["nlist"],
+            quantizer=parameters["quantizer"],
+            nprobe=parameters["nprobe"],
+            train_sample_size=parameters.get("train_sample_size"),
+        ),
+        **parameters,
+    )