Add some query_namespaces performance tweaks

Author: jhamon

poetry.lock

@@ -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"

pyproject.toml

@@ -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)

tests/perf/test_query_namespaces.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)