From 340b62fdf8ddae18062ab785067414f9870b1a95 Mon Sep 17 00:00:00 2001
From: Jen Hamon <jhamon@pinecone.io>
Date: Wed, 13 Nov 2024 10:42:56 -0500
Subject: [PATCH] Add some query_namespaces performance tweaks

---
 poetry.lock                                 | 33 ++++++++++++++-
 pyproject.toml                              |  3 ++
 tests/perf/test_query_namespaces.py         | 45 +++++++++++++++++++++
 tests/perf/test_query_results_aggregator.py | 24 +++++++++++
 4 files changed, 104 insertions(+), 1 deletion(-)
 create mode 100644 tests/perf/test_query_namespaces.py
 create mode 100644 tests/perf/test_query_results_aggregator.py

diff --git a/poetry.lock b/poetry.lock
index 3e918a7e..88303c9f 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -1071,6 +1071,17 @@ files = [
 googleapis-common-protos = "*"
 protobuf = ">=4.21.0"
 
+[[package]]
+name = "py-cpuinfo"
+version = "9.0.0"
+description = "Get CPU info with pure Python"
+optional = false
+python-versions = "*"
+files = [
+    {file = "py-cpuinfo-9.0.0.tar.gz", hash = "sha256:3cdbbf3fac90dc6f118bfd64384f309edeadd902d7c8fb17f02ffa1fc3f49690"},
+    {file = "py_cpuinfo-9.0.0-py3-none-any.whl", hash = "sha256:859625bc251f64e21f077d099d4162689c762b5d6a4c3c97553d56241c9674d5"},
+]
+
 [[package]]
 name = "pygments"
 version = "2.16.1"
@@ -1124,6 +1135,26 @@ pytest = ">=5.4.0"
 [package.extras]
 testing = ["coverage", "hypothesis (>=5.7.1)"]
 
+[[package]]
+name = "pytest-benchmark"
+version = "5.0.0"
+description = "A ``pytest`` fixture for benchmarking code. It will group the tests into rounds that are calibrated to the chosen timer."
+optional = false
+python-versions = ">=3.9"
+files = [
+    {file = "pytest-benchmark-5.0.0.tar.gz", hash = "sha256:cd0adf68516eea7ac212b78a7eb6fc3373865507de8562bb3bfff2f2f852cc63"},
+    {file = "pytest_benchmark-5.0.0-py3-none-any.whl", hash = "sha256:67fed4943aa761077345119555d7f6df09877a12a36e8128f05e19ccd5942d80"},
+]
+
+[package.dependencies]
+py-cpuinfo = "*"
+pytest = ">=3.8"
+
+[package.extras]
+aspect = ["aspectlib"]
+elasticsearch = ["elasticsearch"]
+histogram = ["pygal", "pygaljs", "setuptools"]
+
 [[package]]
 name = "pytest-cov"
 version = "2.10.1"
@@ -1586,4 +1617,4 @@ grpc = ["googleapis-common-protos", "grpcio", "grpcio", "lz4", "protobuf", "prot
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.8"
-content-hash = "d680a8699ebcc13e2369221c4fa987d1206f1a5549393055040fa78dac4da5be"
+content-hash = "0823a7b71260e2e723e281446f1995a1033a3301cbb933b605b202ea55583e4d"
diff --git a/pyproject.toml b/pyproject.toml
index 7c79e3d9..9d170687 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -82,6 +82,9 @@ pytest-asyncio = "0.15.1"
 pytest-cov = "2.10.1"
 pytest-mock = "3.6.1"
 pytest-timeout = "2.2.0"
+pytest-benchmark = [
+  { version = '5.0.0', python = ">=3.9,<4.0" }
+]
 urllib3_mock = "0.3.3"
 responses = ">=0.8.1"
 ddtrace = "^2.14.4"
diff --git a/tests/perf/test_query_namespaces.py b/tests/perf/test_query_namespaces.py
new file mode 100644
index 00000000..22a2492a
--- /dev/null
+++ b/tests/perf/test_query_namespaces.py
@@ -0,0 +1,45 @@
+import time
+import random
+import pytest
+from pinecone import Pinecone
+from pinecone.grpc import PineconeGRPC
+
+latencies = []
+
+
+def call_n_threads(index):
+    query_vec = [random.random() for i in range(1024)]
+    start = time.time()
+    combined_results = index.query_namespaces(
+        vector=query_vec,
+        namespaces=["ns1", "ns2", "ns3", "ns4"],
+        include_values=False,
+        include_metadata=True,
+        filter={"publication_date": {"$eq": "Last3Months"}},
+        top_k=1000,
+    )
+    finish = time.time()
+    # print(f"Query took {finish-start} seconds")
+    latencies.append(finish - start)
+
+    return combined_results
+
+
+class TestQueryNamespacesRest:
+    @pytest.mark.parametrize("n_threads", [4])
+    def test_query_namespaces_grpc(self, benchmark, n_threads):
+        pc = PineconeGRPC()
+        index = pc.Index(
+            host="jen1024-dojoi3u.svc.apw5-4e34-81fa.pinecone.io", pool_threads=n_threads
+        )
+        benchmark.pedantic(call_n_threads, (index,), rounds=10, warmup_rounds=1, iterations=5)
+
+    @pytest.mark.parametrize("n_threads", [4])
+    def test_query_namespaces_rest(self, benchmark, n_threads):
+        pc = Pinecone()
+        index = pc.Index(
+            host="jen1024-dojoi3u.svc.apw5-4e34-81fa.pinecone.io",
+            pool_threads=n_threads,
+            connection_pool_maxsize=20,
+        )
+        benchmark.pedantic(call_n_threads, (index,), rounds=10, warmup_rounds=1, iterations=5)
diff --git a/tests/perf/test_query_results_aggregator.py b/tests/perf/test_query_results_aggregator.py
new file mode 100644
index 00000000..29ac4c35
--- /dev/null
+++ b/tests/perf/test_query_results_aggregator.py
@@ -0,0 +1,24 @@
+import random
+from pinecone.data.query_results_aggregator import QueryResultsAggregator
+
+
+def fake_results(i):
+    matches = [
+        {"id": f"id{i}", "score": random.random(), "values": [random.random() for _ in range(768)]}
+        for _ in range(1000)
+    ]
+    matches.sort(key=lambda x: x["score"], reverse=True)
+    return {"namespace": f"ns{i}", "matches": matches}
+
+
+def aggregate_results(responses):
+    ag = QueryResultsAggregator(1000)
+    for response in responses:
+        ag.add_results(response)
+    return ag.get_results()
+
+
+class TestQueryResultsAggregatorPerf:
+    def test_my_stuff(self, benchmark):
+        responses = [fake_results(i) for i in range(10)]
+        benchmark(aggregate_results, responses)