vectordb-bench 0.0.6__tar.gz → 0.0.8__tar.gz

vectordb_bench-0.0.8/.devcontainer/Dockerfile

@@ -0,0 +1,10 @@
+FROM python:3.11-buster as builder-image
+
+RUN apt-get update
+
+COPY ../install/requirements_py3.11.txt .
+RUN pip3 install -U pip
+RUN pip3 install --no-cache-dir -r requirements_py3.11.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
+
+WORKDIR /opt/code
+ENV PYTHONPATH /opt/code/VectorDBBench

vectordb_bench-0.0.8/.devcontainer/devcontainer.json

@@ -0,0 +1,47 @@
+// For format details, see https://aka.ms/devcontainer.json. For config options, see the
+// README at: https://github.com/devcontainers/templates/tree/main/src/docker-existing-dockerfile
+{
+	"name": "VectorDBBench dev container",
+	"build": {
+		// Sets the run context to one level up instead of the .devcontainer folder.
+		"context": "..",
+		// Update the 'dockerFile' property if you aren't using the standard 'Dockerfile' filename.
+		"dockerfile": "./Dockerfile"
+	},
+	"runArgs": [
+		"--privileged",
+		"--cap-add=SYS_PTRACE"
+	],
+	"mounts": [
+		// You have to make sure source directory is avaliable on your host file system.
+		"source=${localEnv:HOME}/vectordb_bench/dataset,target=/tmp/vectordb_bench/dataset,type=bind,consistency=cached"
+	],
+	"workspaceMount": "source=${localWorkspaceFolder},target=/opt/code/VectorDBBench,type=bind,consistency=cached",
+	"workspaceFolder": "/opt/code/VectorDBBench",
+
+	// Features to add to the dev container. More info: https://containers.dev/features.
+	// "features": {},
+
+	// Use 'forwardPorts' to make a list of ports inside the container available locally.
+	"forwardPorts": [
+		8501
+	],
+
+	// Uncomment the next line to run commands after the container is created.
+	// "postCreateCommand": "cat /etc/os-release",
+
+	// Configure tool-specific properties.
+	"customizations": {
+		"vscode": {
+			"extensions": [
+				"eamodio.gitlens",
+				"ms-python.python",
+				"ms-python.debugpy",
+				"ms-azuretools.vscode-docker"
+			]
+		}
+	}
+
+	// Uncomment to connect as an existing user other than the container default. More info: https://aka.ms/dev-containers-non-root.
+	// "remoteUser": "devcontainer"
+}

vectordb_bench-0.0.8/.github/workflows/pull_request.yml

@@ -0,0 +1,36 @@
+name: Test on pull request
+
+on:
+  pull_request:
+    branches:
+      - main
+
+jobs:
+  build:
+    name: Run Python Tests
+    strategy:
+      matrix:
+        python-version: [3.11, 3.12]
+        os: [ubuntu-latest, windows-latest]
+    runs-on: ${{ matrix.os }}
+
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+      - name: Setup Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Fetch tags
+        run: |
+          git fetch --prune --unshallow --tags
+
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install -e ".[test]"
+
+      - name: Test with pytest
+        run: |
+          make unittest

vectordb_bench-0.0.8/Dockerfile

@@ -0,0 +1,18 @@
+FROM python:3.11-buster as builder-image
+
+RUN apt-get update
+
+COPY install/requirements_py3.11.txt .
+RUN pip3 install -U pip
+RUN pip3 install --no-cache-dir -r requirements_py3.11.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
+
+FROM python:3.11-slim-buster
+
+COPY --from=builder-image /usr/local/bin /usr/local/bin
+COPY --from=builder-image /usr/local/lib/python3.11/site-packages /usr/local/lib/python3.11/site-packages
+
+WORKDIR /opt/code
+COPY . .
+ENV PYTHONPATH /opt/code
+
+ENTRYPOINT ["python3", "-m", "vectordb_bench"]

vectordb_bench-0.0.8/Makefile

@@ -0,0 +1,2 @@
+unittest:
+	PYTHONPATH=`pwd` python3 -m pytest tests/test_dataset.py::TestDataSet::test_download_small -svv

{vectordb-bench-0.0.6 → vectordb_bench-0.0.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vectordb-bench
-Version: 0.0.6
+Version: 0.0.8
 Summary: 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.
 Author-email: XuanYang-cn <xuan.yang@zilliz.com>
 Project-URL: repository, https://github.com/zilliztech/VectorDBBench
@@ -49,7 +49,7 @@ Provides-Extra: elastic
 Requires-Dist: elasticsearch; extra == "elastic"
 Provides-Extra: pgvector
 Requires-Dist: pgvector; extra == "pgvector"
-Requires-Dist: sqlalchemy; extra == "pgvector"
+Requires-Dist: psycopg2; extra == "pgvector"
 Provides-Extra: pgvecto-rs
 Requires-Dist: psycopg2; extra == "pgvecto-rs"
 Provides-Extra: redis
@@ -149,6 +149,19 @@ OR:
 ```shell
 $ init_bench
 ```
+
+OR:
+
+If you are using [dev container](https://code.visualstudio.com/docs/devcontainers/containers), create
+the following dataset directory first:
+
+```shell
+# Mount local ~/vectordb_bench/dataset to contain's /tmp/vectordb_bench/dataset.
+# If you are not comfortable with the path name, feel free to change it in devcontainer.json
+mkdir -p ~/vectordb_bench/dataset
+```
+After reopen the repository in container, run `python -m vectordb_bench` in the container's bash.
+
 ### Check coding styles
 ```shell
 $ ruff check vectordb_bench

{vectordb-bench-0.0.6 → vectordb_bench-0.0.8}/README.md

@@ -90,6 +90,19 @@ OR:
 ```shell
 $ init_bench
 ```
+
+OR:
+
+If you are using [dev container](https://code.visualstudio.com/docs/devcontainers/containers), create
+the following dataset directory first:
+
+```shell
+# Mount local ~/vectordb_bench/dataset to contain's /tmp/vectordb_bench/dataset.
+# If you are not comfortable with the path name, feel free to change it in devcontainer.json
+mkdir -p ~/vectordb_bench/dataset
+```
+After reopen the repository in container, run `python -m vectordb_bench` in the container's bash.
+
 ### Check coding styles
 ```shell
 $ ruff check vectordb_bench

vectordb_bench-0.0.8/install/requirements_py3.11.txt

@@ -0,0 +1,23 @@
+grpcio==1.53.0
+grpcio-tools==1.53.0
+qdrant-client
+pinecone-client
+weaviate-client
+elasticsearch
+pgvector
+sqlalchemy
+redis
+chromadb
+pytz
+streamlit-autorefresh
+streamlit>=1.23.0
+streamlit_extras
+tqdm
+s3fs
+psutil
+polars
+plotly
+environs
+pydantic<v2
+scikit-learn
+pymilvus

vectordb_bench-0.0.8/install.py

@@ -0,0 +1,72 @@
+import os
+import argparse
+import subprocess
+
+def docker_tag_base():
+    return 'vdbbench'
+
+def dockerfile_path_base():
+    return os.path.join('vectordb_bench/', '../Dockerfile')
+
+def docker_tag(track, algo):
+    return docker_tag_base() + '-' + track + '-' + algo
+
+
+def build(tag, args, dockerfile):
+    print('Building %s...' % tag)
+    if args is not None and len(args) != 0:
+        q = " ".join(["--build-arg " + x.replace(" ", "\\ ") for x in args])
+    else:
+        q = ""
+
+    try:
+        command = 'docker build %s --rm -t %s -f' \
+                   % (q, tag)
+        command += ' %s .' % dockerfile
+        print(command)
+        subprocess.check_call(command, shell=True)
+        return {tag: 'success'}
+    except subprocess.CalledProcessError:
+        return {tag: 'fail'}
+
+def build_multiprocess(args):
+    return build(*args)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    parser.add_argument(
+        "--proc",
+        default=1,
+        type=int,
+        help="the number of process to build docker images")
+    parser.add_argument(
+        '--track',
+        choices=['none'],
+        default='none'
+    )
+    parser.add_argument(
+        '--algorithm',
+        metavar='NAME',
+        help='build only the named algorithm image',
+        default=None)
+    parser.add_argument(
+        '--dockerfile',
+        metavar='PATH',
+        help='build only the image from a Dockerfile path',
+        default=None)
+    parser.add_argument(
+        '--build-arg',
+        help='pass given args to all docker builds',
+        nargs="+")
+    args = parser.parse_args()
+
+    print('Building base image...')
+
+    subprocess.check_call(
+        'docker build \
+        --rm -t %s -f %s .' % (docker_tag_base(), dockerfile_path_base()), shell=True)
+
+    print('Building end.')
+

{vectordb-bench-0.0.6 → vectordb_bench-0.0.8}/pyproject.toml

@@ -5,6 +5,10 @@ build-backend = "setuptools.build_meta"
 [tool.setuptools.package-data]
 "vectordb_bench.results" = ["*.json"]
 
+[tool.setuptools.packages.find]
+where = ["."]
+include = ["vectordb_bench"]
+
 [project]
 name = "vectordb-bench"
 authors = [
@@ -61,7 +65,7 @@ qdrant = [ "qdrant-client" ]
 pinecone = [ "pinecone-client" ]
 weaviate = [ "weaviate-client" ]
 elastic = [ "elasticsearch" ]
-pgvector = [ "pgvector", "sqlalchemy" ]
+pgvector = [ "pgvector", "psycopg2" ]
 pgvecto_rs = [ "psycopg2" ]
 redis = [ "redis" ]
 chromadb = [ "chromadb" ]

{vectordb-bench-0.0.6 → vectordb_bench-0.0.8}/tests/test_chroma.py

@@ -52,6 +52,7 @@ class TestChroma:
 
         # insert
         with chrma.init():
+            #chrma.client.delete_collection("example2")
             assert (chrma.client.heartbeat() is not None), "chroma client is not connected"
             res = chrma.insert_embeddings(embeddings=embeddings, metadata=range(count))
             # bulk_insert return
@@ -87,7 +88,7 @@ class TestChroma:
 
 
             res = chrma.search_embedding(
-                query=q, k=100, filters={"metadata": filter_value}
+                query=q, k=100, filters={"id": filter_value}
             )
             assert (
                 res[0] == int(test_id)
@@ -101,18 +102,4 @@ class TestChroma:
                     break
             assert isFilter, f"Filter not working, id_list: {id_list}"
 
-            #Test id filter
-            res = chrma.search_embedding(
-                query=q, k=100, filters={"id": 9999}
-            )
-            assert (
-                res[0] == 9999
-            )
-
-            #Test two filters, id and metadata
-            res = chrma.search_embedding(
-                query=q, k=100, filters={"metadata": filter_value, "id": 9999}
-            )
-            assert (
-                res[0] == 9999 and len(res) == 1, f"filters failed, got: ({res[0]}), expected ({9999})"
-            )
+            

vectordb_bench-0.0.8/tests/test_data_source.py

@@ -0,0 +1,28 @@
+import logging
+import pytest
+from vectordb_bench.backend.data_source import DatasetSource
+from vectordb_bench.backend.cases import type2case
+
+log = logging.getLogger("vectordb_bench")
+
+class TestReader:
+    @pytest.mark.parametrize("type_case", [
+        (k, v) for k, v in type2case.items()
+    ])
+    def test_type_cases(self, type_case):
+        self.per_case_test(type_case)
+
+
+    def per_case_test(self, type_case):
+        t, ca_cls = type_case
+        ca = ca_cls()
+        log.info(f"test case: {t.name}, {ca.name}")
+
+        filters = ca.filter_rate
+        ca.dataset.prepare(source=DatasetSource.AliyunOSS, filters=filters)
+        ali_trains = ca.dataset.train_files
+
+        ca.dataset.prepare(filters=filters)
+        s3_trains = ca.dataset.train_files
+
+        assert ali_trains == s3_trains

vectordb_bench-0.0.8/tests/test_dataset.py

@@ -0,0 +1,77 @@
+from vectordb_bench.backend.dataset import Dataset
+import logging
+import pytest
+from pydantic import ValidationError
+from vectordb_bench.backend.data_source import DatasetSource
+
+
+log = logging.getLogger("vectordb_bench")
+
+class TestDataSet:
+    def test_iter_dataset(self):
+        for ds in Dataset:
+            log.info(ds)
+
+    def test_cohere(self):
+        cohere = Dataset.COHERE.get(100_000)
+        log.info(cohere)
+        assert cohere.name == "Cohere"
+        assert cohere.size == 100_000
+        assert cohere.label == "SMALL"
+        assert cohere.dim == 768
+
+    def test_cohere_error(self):
+        with pytest.raises(ValidationError):
+            Dataset.COHERE.get(9999)
+
+    def test_iter_cohere(self):
+        cohere_10m = Dataset.COHERE.manager(10_000_000)
+        cohere_10m.prepare()
+
+        import time
+        before = time.time()
+        for i in cohere_10m:
+            log.debug(i.head(1))
+
+        dur_iter = time.time() - before
+        log.warning(f"iter through cohere_10m cost={dur_iter/60}min")
+
+    # pytest -sv tests/test_dataset.py::TestDataSet::test_iter_laion
+    def test_iter_laion(self):
+        laion_100m = Dataset.LAION.manager(100_000_000)
+        from vectordb_bench.backend.data_source import DatasetSource
+        laion_100m.prepare(source=DatasetSource.AliyunOSS)
+
+        import time
+        before = time.time()
+        for i in laion_100m:
+            log.debug(i.head(1))
+
+        dur_iter = time.time() - before
+        log.warning(f"iter through laion_100m cost={dur_iter/60}min")
+
+    def test_download_small(self):
+        openai_50k = Dataset.OPENAI.manager(50_000)
+        files = [
+            "test.parquet",
+            "neighbors.parquet",
+            "neighbors_head_1p.parquet",
+            "neighbors_tail_1p.parquet",
+        ]
+
+        file_path = openai_50k.data_dir.joinpath("test.parquet")
+        import os
+
+        DatasetSource.S3.reader().read(
+            openai_50k.data.dir_name.lower(),
+            files=files,
+            local_ds_root=openai_50k.data_dir,
+        )
+
+        os.remove(file_path)
+        DatasetSource.AliyunOSS.reader().read(
+            openai_50k.data.dir_name.lower(),
+            files=files,
+            local_ds_root=openai_50k.data_dir,
+        )
+

{vectordb-bench-0.0.6 → vectordb_bench-0.0.8}/tests/test_utils.py

@@ -37,3 +37,30 @@ class TestUtils:
         log.info(f"recall: {res}, expected: {expected}")
         assert res == expected
 
+
+class TestGetFiles:
+    @pytest.mark.parametrize("train_count", [
+        1,
+        10,
+        50,
+        100,
+    ])
+    def test_train_count(self, train_count):
+        files = utils.compose_train_files(train_count, True)
+        log.info(files)
+
+        assert len(files) == train_count
+
+    @pytest.mark.parametrize("use_shuffled", [True, False])
+    def test_use_shuffled(self, use_shuffled):
+        files = utils.compose_train_files(1, use_shuffled)
+        log.info(files)
+
+        trains = [f for f in files if "train" in f]
+        if use_shuffled:
+            for t in trains:
+                assert "shuffle_train" in t
+        else:
+            for t in trains:
+                assert "shuffle" not in t
+                assert "train" in t

{vectordb-bench-0.0.6 → vectordb_bench-0.0.8}/vectordb_bench/backend/clients/__init__.py

@@ -54,8 +54,8 @@ class DB(Enum):
             return ElasticCloud
 
         if self == DB.QdrantCloud:
-            from .qdrant_cloud.qdrant_cloud import QdrantClient
-            return QdrantClient
+            from .qdrant_cloud.qdrant_cloud import QdrantCloud
+            return QdrantCloud
 
         if self == DB.WeaviateCloud:
             from .weaviate_cloud.weaviate_cloud import WeaviateCloud
@@ -142,8 +142,8 @@ class DB(Enum):
             return WeaviateIndexConfig
 
         if self == DB.PgVector:
-            from .pgvector.config import PgVectorIndexConfig
-            return PgVectorIndexConfig
+            from .pgvector.config import _pgvector_case_config
+            return _pgvector_case_config.get(index_type)
 
         if self == DB.PgVectoRS:
             from .pgvecto_rs.config import _pgvecto_rs_case_config

{vectordb-bench-0.0.6 → vectordb_bench-0.0.8}/vectordb_bench/backend/clients/api.py

@@ -16,6 +16,7 @@ class IndexType(str, Enum):
     HNSW = "HNSW"
     DISKANN = "DISKANN"
     IVFFlat = "IVF_FLAT"
+    IVFSQ8 = "IVF_SQ8"
     Flat = "FLAT"
     AUTOINDEX = "AUTOINDEX"
     ES_HNSW = "hnsw"

{vectordb-bench-0.0.6 → vectordb_bench-0.0.8}/vectordb_bench/backend/clients/chroma/chroma.py

@@ -106,21 +106,9 @@ class ChromaClient(VectorDB):
         """
         if filters:
             # assumes benchmark test filters of format: {'metadata': '>=10000', 'id': 10000}
-            metadata_value = filters.get("metadata")
             id_value = filters.get("id")
-            if metadata_value and id_value:
-                results = self.collection.query(
-                    query_embeddings=query, n_results=k, 
-                    where={"$and": [{"id": {"$eq": id_value}}, 
-                                    {"id": {"$gt": metadata_value}}
-                                    ]}
-                )
-            elif metadata_value:
-                results = self.collection.query(query_embeddings=query, n_results=k, 
-                                                where={"id": {"$gt": metadata_value}})
-            else:
-                results = self.collection.query(query_embeddings=query, n_results=k, 
-                                                where={"id": {"$eq": id_value}})
+            results = self.collection.query(query_embeddings=query, n_results=k, 
+                                                where={"id": {"$gt": id_value}})
             #return list of id's in results
             return [int(i) for i in results.get('ids')[0]]
         results = self.collection.query(query_embeddings=query, n_results=k)

{vectordb-bench-0.0.6 → vectordb_bench-0.0.8}/vectordb_bench/backend/clients/milvus/config.py

@@ -95,6 +95,24 @@ class IVFFlatConfig(MilvusIndexConfig, DBCaseConfig):
             "metric_type": self.parse_metric(),
             "params": {"nprobe": self.nprobe},
         }
+        
+class IVFSQ8Config(MilvusIndexConfig, DBCaseConfig):
+    nlist: int
+    nprobe: int | None = None
+    index: IndexType = IndexType.IVFSQ8
+
+    def index_param(self) -> dict:
+        return {
+            "metric_type": self.parse_metric(),
+            "index_type": self.index.value,
+            "params": {"nlist": self.nlist},
+        }
+
+    def search_param(self) -> dict:
+        return {
+            "metric_type": self.parse_metric(),
+            "params": {"nprobe": self.nprobe},
+        }
 
 
 class FLATConfig(MilvusIndexConfig, DBCaseConfig):
@@ -210,6 +228,7 @@ _milvus_case_config = {
     IndexType.HNSW: HNSWConfig,
     IndexType.DISKANN: DISKANNConfig,
     IndexType.IVFFlat: IVFFlatConfig,
+    IndexType.IVFSQ8: IVFSQ8Config,
     IndexType.Flat: FLATConfig,
     IndexType.GPU_IVF_FLAT: GPUIVFFlatConfig,
     IndexType.GPU_IVF_PQ: GPUIVFPQConfig,