Add support for OpenSearch as a database (#300)

## Problem

Describe the purpose of this change. What problem is being solved and
why?

## Solution

Describe the approach you took. Link to any relevant bugs, issues, docs,
or other resources.

## Type of Change

- [ ] Bug fix (non-breaking change which fixes an issue)
- [ ] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to not work as expected)
- [ ] This change requires a documentation update
- [ ] Infrastructure change (CI configs, etc)
- [ ] Non-code change (docs, etc)
- [ ] None of the above: (explain here)

## Test Plan

Describe specific steps for validating this change.

---------

Co-authored-by: Vamshi Krishna Enabothala <[email protected]>

Author: daverigby

Diff for: .github/workflows/python-package.yml

@@ -19,7 +19,7 @@ jobs:
   build:
 
     runs-on: ubuntu-latest
-    timeout-minutes: 15
+    timeout-minutes: 30
     strategy:
       fail-fast: false
       matrix:
@@ -52,6 +52,7 @@ jobs:
     - name: Set up Docker Compose
       run: |
         docker compose -f docker/pgvector/docker-compose.yml up -d
+        docker compose -f docker/opensearch/docker-compose.yml up -d
     - name: Test with pytest
       run: |
         # Maximum observed test runtime in CI is ~60s. Set a per-test timeout of
@@ -76,3 +77,4 @@ jobs:
       if: always()
       run: |
         docker compose -f docker/pgvector/docker-compose.yml down
+        docker compose -f docker/opensearch/docker-compose.yml down

Diff for: .gitignore

@@ -158,3 +158,9 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 .idea/
+
+# reports
+reports/
+
+# DS_Store
+.DS_Store

Diff for: README.md

@@ -110,6 +110,7 @@ The following databases are currently supported by VSB:
 
 * [Pinecone](vsb/databases/pinecone/README.md)
 * [pgvector](vsb/databases/pgvector/README.md)
+* [OpenSearch](vsb/databases/opensearch/README.md)
 
 > [!TIP]
 > You can also display the list of supported databases using the following command: 

Diff for: docker/opensearch/docker-compose.yml

@@ -0,0 +1,22 @@
+services:
+  opensearch:
+    image: opensearchproject/opensearch:2
+    container_name: vsb_opensearch
+    environment:
+      discovery.type: single-node
+      # Necessary to be able to use non-TLS (for simplifying local testing)
+      DISABLE_SECURITY_PLUGIN: true
+    deploy:
+      resources:
+        # By default, limit to 8GB RAM & 2 CPU cores - reasonable values which
+        # should be sufficient for a range of workloads and typical commodity
+        # hardware can support.
+        # Adjust as needed for larger workloads.
+        limits:
+          memory: 8GB
+          cpus: '2'
+    shm_size: 8GB # Equal to limits.memory.
+    ports:
+      - "9200:9200"
+      - "9600:9600"
+    restart: always

Diff for: poetry.lock

@@ -26,6 +26,8 @@ hdrhistogram = "^0.10.3"
 tenacity = "^9.0.0"
 rich = "^13.8.1"
 filelock = "^3.18.0"
+opensearch-py = "^2.8.0"
+requests-aws4auth = "^1.3.1"
 
 [tool.poetry.scripts]
 vsb = "vsb.main:main"

Diff for: pyproject.toml

@@ -26,9 +26,12 @@
     spawn_vsb_pinecone,
 )
 from test_pgvector import spawn_vsb_pgvector
+from test_opensearch import spawn_vsb_opensearch
 
 
-@pytest.mark.parametrize("spawn_vsb", [spawn_vsb_pgvector, spawn_vsb_pinecone])
+@pytest.mark.parametrize(
+    "spawn_vsb", [spawn_vsb_pgvector, spawn_vsb_pinecone, spawn_vsb_opensearch]
+)
 class TestCommon:
 
     # Unfortunately pytest won't let us selectively parametrize with fixtures, so
@@ -52,8 +55,8 @@ def test_mnist_single(
         check_request_counts(
             stdout,
             {
-                # Populate num_requests counts batches, not individual records.
-                "Populate": {"num_requests": lambda x: x <= 2, "num_failures": 0},
+                # Populate num_requests counts batches, not individual records (600).
+                "Populate": {"num_requests": lambda x: x < 600, "num_failures": 0},
                 "Search": {
                     "num_requests": 20,
                     "num_failures": 0,
@@ -82,10 +85,10 @@ def test_mnist_concurrent(
             stdout,
             {
                 # For multiple users the populate phase will chunk the records to be
-                # loaded into num_users chunks - i.e. 4 here. Given the size of each
-                # chunk will be less than the batch size (600 / 4 < 1000), then the
-                # number of requests will be equal to the number of users - i.e. 4
-                "Populate": {"num_requests": 4, "num_failures": 0},
+                # loaded into num_users chunks - i.e. 4 here. Different DBs
+                # use different batch sizes, so just check we have fewer than
+                # number of records (600) / number of users (4).
+                "Populate": {"num_requests": lambda x: x < 600 / 4, "num_failures": 0},
                 "Search": {
                     "num_requests": 20,
                     "num_failures": 0,
@@ -114,10 +117,10 @@ def test_mnist_multiprocess(
             stdout,
             {
                 # For multiple users the populate phase will chunk the records to be
-                # loaded into num_users chunks - i.e. 4 here. Given the size of each
-                # chunk will be less than the batch size (600 / 4 < 1000), then the
-                # number of requests will be equal to the number of users - i.e. 4
-                "Populate": {"num_requests": 4, "num_failures": 0},
+                # loaded into num_users chunks - i.e. 4 here. Different DBs
+                # use different batch sizes, so just check we have fewer than
+                # number of records (600) / number of users (4).
+                "Populate": {"num_requests": lambda x: x < 600 / 4, "num_failures": 0},
                 # The number of Search requests should equal the number in the dataset
                 # (20 for mnist-test).
                 "Search": {
@@ -145,15 +148,22 @@ def test_mnist_double(
         check_request_counts(
             stdout,
             {
-                "test1.Populate": {"num_requests": lambda x: x <= 2, "num_failures": 0},
+                # Populate num_requests counts batches, not individual records (600).
+                "test1.Populate": {
+                    "num_requests": lambda x: x < 600,
+                    "num_failures": 0,
+                },
                 # The number of Search requests should equal the number in the dataset
                 # (20 for mnist-test).
                 "test1.Search": {
                     "num_requests": 20,
                     "num_failures": 0,
                     "Recall": check_recall_stats,
                 },
-                "test2.Populate": {"num_requests": lambda x: x <= 2, "num_failures": 0},
+                "test2.Populate": {
+                    "num_requests": lambda x: x < 600,
+                    "num_failures": 0,
+                },
                 "test2.Search": {
                     "num_requests": 20,
                     "num_failures": 0,
@@ -182,18 +192,24 @@ def test_mnist_double_concurrent(
             stdout,
             {
                 # For multiple users the populate phase will chunk the records to be
-                # loaded into num_users chunks - i.e. 4 here. Given the size of each
-                # chunk will be less than the batch size (600 / 4 < 200), then the
-                # number of requests will be equal to the number of users - i.e. 4
-                "test1.Populate": {"num_requests": 4, "num_failures": 0},
+                # loaded into num_users chunks - i.e. 4 here. Different DBs
+                # use different batch sizes, so just check we have fewer than
+                # number of records (600) / number of users (4).
+                "test1.Populate": {
+                    "num_requests": lambda x: x < 600 / 4,
+                    "num_failures": 0,
+                },
                 # The number of Search requests should equal the number in the dataset
                 # (20 for mnist-test).
                 "test1.Search": {
                     "num_requests": 20,
                     "num_failures": 0,
                     "Recall": check_recall_stats,
                 },
-                "test2.Populate": {"num_requests": 4, "num_failures": 0},
+                "test2.Populate": {
+                    "num_requests": lambda x: x < 600 / 4,
+                    "num_failures": 0,
+                },
                 "test2.Search": {
                     "num_requests": 20,
                     "num_failures": 0,
@@ -222,18 +238,28 @@ def test_mnist_double_multiprocess(
             stdout,
             {
                 # For multiple users the populate phase will chunk the records to be
-                # loaded into num_users chunks - i.e. 4 here. Given the size of each
-                # chunk will be less than the batch size (600 / 4 < 200), then the
-                # number of requests will be equal to the number of users - i.e. 4
-                "test1.Populate": {"num_requests": 4, "num_failures": 0},
+                # loaded into num_users chunks - i.e. 4 here. Different DBs
+                # use different batch sizes, so just check we have fewer than
+                # number of records (600) / number of users (4).
+                "test1.Populate": {
+                    "num_requests": lambda x: x < 600 / 4,
+                    "num_failures": 0,
+                },
                 # The number of Search requests should equal the number in the dataset
                 # (20 for mnist-test).
                 "test1.Search": {
                     "num_requests": 20,
                     "num_failures": 0,
                     "Recall": check_recall_stats,
                 },
-                "test2.Populate": {"num_requests": 4, "num_failures": 0},
+                # For multiple users the populate phase will chunk the records to be
+                # loaded into num_users chunks - i.e. 4 here. Different DBs
+                # use different batch sizes, so just check we have fewer than
+                # number of records (600) / number of users (4).
+                "test2.Populate": {
+                    "num_requests": lambda x: x < 600 / 4,
+                    "num_failures": 0,
+                },
                 "test2.Search": {
                     "num_requests": 20,
                     "num_failures": 0,
@@ -262,8 +288,8 @@ def test_mnist_skip_populate(
         check_request_counts(
             stdout,
             {
-                # Populate num_requests counts batches, not individual records.
-                "Populate": {"num_requests": lambda x: x <= 2, "num_failures": 0},
+                # Populate num_requests counts batches, not individual records (600).
+                "Populate": {"num_requests": lambda x: x < 600, "num_failures": 0},
                 "Search": {"num_requests": 20, "num_failures": 0},
             },
         )
@@ -307,7 +333,7 @@ def test_filtered(
             {
                 # Populate num_requests counts batches, not individual records.
                 "Populate": {
-                    "num_requests": lambda x: x == 10 or x == 210,
+                    "num_requests": lambda x: x > 1 and x < 10000,
                     "num_failures": 0,
                 },
                 "Search": {
@@ -335,7 +361,7 @@ def test_synthetic(
         check_request_counts(
             stdout,
             {
-                "Populate": {"num_requests": 10, "num_failures": 0},
+                "Populate": {"num_failures": 0},
                 "Search": {
                     "num_requests": 100,
                     "num_failures": 0,
@@ -367,7 +393,7 @@ def test_synthetic_runbook(
         check_request_counts(
             stdout,
             {
-                "Populate": {"num_requests": lambda x: x <= 4, "num_failures": 0},
+                "Populate": {"num_failures": 0},
                 "Search": {
                     "num_requests": 500,
                     "num_failures": 0,
@@ -382,6 +408,12 @@ def test_synthetic_proportional(
         pinecone_api_key,
         pinecone_index_synthetic,
     ):
+        if spawn_vsb == spawn_vsb_opensearch:
+            pytest.skip(
+                "Synthetic proportional test not supported on OpenSearch ("
+                "fetch_batch not yet implemented for OpenSearch)"
+            )
+
         (proc, stdout, stderr) = spawn_vsb(
             pinecone_api_key=pinecone_api_key,
             pinecone_index=pinecone_index_synthetic,
@@ -404,7 +436,7 @@ def test_synthetic_proportional(
         check_request_counts(
             stdout,
             {
-                "Populate": {"num_requests": lambda x: x <= 4, "num_failures": 0},
+                "Populate": {"num_failures": 0},
                 "Search": {
                     "num_requests": lambda x: (x >= 150 and x <= 250),
                     "num_failures": 0,

Diff for: tests/integration/test_common.py

@@ -0,0 +1,16 @@
+from conftest import (
+    spawn_vsb_inner,
+)
+
+
+# used in test_common
+def spawn_vsb_opensearch(workload, timeout=120, extra_args=None, **kwargs):
+    """Spawn an instance of pgvector vsb with the given arguments, returning the proc object,
+    its stdout and stderr.
+    """
+    extra_env = {
+        "VSB__OPENSEARCH_USERNAME": "admin",
+        "VSB__OPENSEARCH_PASSWORD": "opensearch",
+        "VSB__OPENSEARCH_USE_TLS": "false",
+    }
+    return spawn_vsb_inner("opensearch", workload, timeout, extra_args, extra_env)