Skip to content

Commit 8399088

Browse files
SK1PPRSK1PPR
authored
Benchmarking optimizations (#84)
* Reduce in-memory checkpoint lock scope * Optimize HNSW distance calculations * Expose HNSW tuning parameters * Add Python batch vector client methods * Expose KD-tree rebuild tuning parameters
1 parent e49d4c2 commit 8399088

18 files changed

Lines changed: 467 additions & 68 deletions

File tree

‎client/python/tests/test_client.py‎

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,34 @@ def test_insert_rejects_invalid_vector(client):
5454
)
5555

5656

57+
def test_insert_batch_success(client, mock_connection):
58+
mock_connection.call.return_value = Mock(
59+
ids=[
60+
Mock(id=Mock(value="p1")),
61+
Mock(id=Mock(value="p2")),
62+
]
63+
)
64+
65+
point_ids = client.insert_batch(
66+
points=[
67+
(DenseVector([1, 2, 3]), Payload.text("hello")),
68+
(DenseVector([4, 5, 6]), Payload.text("world")),
69+
]
70+
)
71+
72+
assert point_ids == ["p1", "p2"]
73+
74+
75+
def test_insert_batch_rejects_invalid_vector(client):
76+
with pytest.raises(TypeError):
77+
client.insert_batch(
78+
points=[
79+
(DenseVector([1, 2, 3]), Payload.text("hello")),
80+
([4, 5, 6], Payload.text("world")),
81+
]
82+
)
83+
84+
5785
# Get
5886

5987
def test_get_point_success(client, mock_connection):
@@ -118,6 +146,43 @@ def test_search_invalid_vector(client):
118146
)
119147

120148

149+
def test_search_batch_success(client, mock_connection):
150+
mock_connection.call.return_value = Mock(
151+
results=[
152+
Mock(
153+
result_point_ids=[
154+
Mock(id=Mock(value="p1")),
155+
Mock(id=Mock(value="p2")),
156+
]
157+
),
158+
Mock(
159+
result_point_ids=[
160+
Mock(id=Mock(value="p3")),
161+
]
162+
),
163+
]
164+
)
165+
166+
results = client.search_batch(
167+
queries=[
168+
(DenseVector([1, 2, 3]), Similarity.COSINE, 2),
169+
(DenseVector([4, 5, 6]), Similarity.COSINE, 1),
170+
]
171+
)
172+
173+
assert results == [["p1", "p2"], ["p3"]]
174+
175+
176+
def test_search_batch_rejects_invalid_vector(client):
177+
with pytest.raises(TypeError):
178+
client.search_batch(
179+
queries=[
180+
(DenseVector([1, 2, 3]), Similarity.COSINE, 2),
181+
([4, 5, 6], Similarity.COSINE, 1),
182+
]
183+
)
184+
185+
121186
# Close
122187

123188
def test_close_closes_connection(client, mock_connection):

‎client/python/vortexdb/client.py‎

Lines changed: 60 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
from typing import List
1+
from typing import List, Sequence
22

33
from vortexdb.connection import GRPCConnection
44
from vortexdb.config import VortexDBConfig
@@ -38,11 +38,7 @@ def insert(self, *, vector: DenseVector, payload: Payload) -> str:
3838
Insert a vector with payload.
3939
Returns: point_id (str)
4040
"""
41-
if not isinstance(vector, DenseVector):
42-
raise TypeError(
43-
"vector must be a DenseVector. "
44-
"Use: DenseVector([1.0, 2.0, 3.0])"
45-
)
41+
self._validate_dense_vector(vector)
4642

4743
request = proto.build_insert_request(
4844
vector=vector,
@@ -56,6 +52,29 @@ def insert(self, *, vector: DenseVector, payload: Payload) -> str:
5652

5753
return response.id.value
5854

55+
def insert_batch(
56+
self,
57+
*,
58+
points: Sequence[tuple[DenseVector, Payload]],
59+
) -> List[str]:
60+
"""
61+
Insert multiple vectors with payloads.
62+
Returns: List of point IDs
63+
"""
64+
for vector, _ in points:
65+
self._validate_dense_vector(vector)
66+
67+
request = proto.build_batch_insert_request(
68+
points=list(points),
69+
)
70+
71+
response = self._conn.call(
72+
self._conn.stub.InsertVectorsBatch,
73+
request,
74+
)
75+
76+
return [pid.id.value for pid in response.ids]
77+
5978
def get(self, *, point_id: str) -> Point | None:
6079
"""
6180
Retrieve a point by ID.
@@ -95,11 +114,7 @@ def search(
95114
Search for nearest neighbors.
96115
Returns: List of point IDs
97116
"""
98-
if not isinstance(vector, DenseVector):
99-
raise TypeError(
100-
"vector must be a DenseVector. "
101-
"Use: DenseVector([1.0, 2.0, 3.0])"
102-
)
117+
self._validate_dense_vector(vector)
103118

104119
request = proto.build_search_request(
105120
vector=vector,
@@ -114,6 +129,40 @@ def search(
114129

115130
return [pid.id.value for pid in response.result_point_ids]
116131

132+
def search_batch(
133+
self,
134+
*,
135+
queries: Sequence[tuple[DenseVector, Similarity, int]],
136+
) -> List[List[str]]:
137+
"""
138+
Search nearest neighbors for multiple query vectors.
139+
Returns: List of result point ID lists
140+
"""
141+
for vector, _, _ in queries:
142+
self._validate_dense_vector(vector)
143+
144+
request = proto.build_batch_search_request(
145+
queries=list(queries),
146+
)
147+
148+
response = self._conn.call(
149+
self._conn.stub.SearchPointsBatch,
150+
request,
151+
)
152+
153+
return [
154+
[pid.id.value for pid in result.result_point_ids]
155+
for result in response.results
156+
]
157+
158+
@staticmethod
159+
def _validate_dense_vector(vector: DenseVector) -> None:
160+
if not isinstance(vector, DenseVector):
161+
raise TypeError(
162+
"vector must be a DenseVector. "
163+
"Use: DenseVector([1.0, 2.0, 3.0])"
164+
)
165+
117166
def close(self) -> None:
118167
"""
119168
Close the gRPC connection.

‎client/python/vortexdb/grpc/vector_db_pb2.py‎

Lines changed: 18 additions & 10 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎client/python/vortexdb/grpc/vector_db_pb2_grpc.py‎

Lines changed: 91 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,9 @@
44
import warnings
55

66
from google.protobuf import empty_pb2 as google_dot_protobuf_dot_empty__pb2
7-
from vortexdb.grpc import vector_db_pb2 as vector__db__pb2
7+
from . import vector_db_pb2 as vector__db__pb2
88

9-
GRPC_GENERATED_VERSION = '1.76.0'
9+
GRPC_GENERATED_VERSION = '1.81.1'
1010
GRPC_VERSION = grpc.__version__
1111
_version_not_supported = False
1212

@@ -26,7 +26,7 @@
2626
)
2727

2828

29-
class VectorDBStub(object):
29+
class VectorDBStub:
3030
"""Missing associated documentation comment in .proto file."""
3131

3232
def __init__(self, channel):
@@ -55,9 +55,19 @@ def __init__(self, channel):
5555
request_serializer=vector__db__pb2.SearchRequest.SerializeToString,
5656
response_deserializer=vector__db__pb2.SearchResponse.FromString,
5757
_registered_method=True)
58+
self.InsertVectorsBatch = channel.unary_unary(
59+
'/vectordb.VectorDB/InsertVectorsBatch',
60+
request_serializer=vector__db__pb2.InsertVectorsBatchRequest.SerializeToString,
61+
response_deserializer=vector__db__pb2.InsertVectorsBatchResponse.FromString,
62+
_registered_method=True)
63+
self.SearchPointsBatch = channel.unary_unary(
64+
'/vectordb.VectorDB/SearchPointsBatch',
65+
request_serializer=vector__db__pb2.SearchPointsBatchRequest.SerializeToString,
66+
response_deserializer=vector__db__pb2.SearchPointsBatchResponse.FromString,
67+
_registered_method=True)
5868

5969

60-
class VectorDBServicer(object):
70+
class VectorDBServicer:
6171
"""Missing associated documentation comment in .proto file."""
6272

6373
def InsertVector(self, request, context):
@@ -88,6 +98,18 @@ def SearchPoints(self, request, context):
8898
context.set_details('Method not implemented!')
8999
raise NotImplementedError('Method not implemented!')
90100

101+
def InsertVectorsBatch(self, request, context):
102+
"""Missing associated documentation comment in .proto file."""
103+
context.set_code(grpc.StatusCode.UNIMPLEMENTED)
104+
context.set_details('Method not implemented!')
105+
raise NotImplementedError('Method not implemented!')
106+
107+
def SearchPointsBatch(self, request, context):
108+
"""Missing associated documentation comment in .proto file."""
109+
context.set_code(grpc.StatusCode.UNIMPLEMENTED)
110+
context.set_details('Method not implemented!')
111+
raise NotImplementedError('Method not implemented!')
112+
91113

92114
def add_VectorDBServicer_to_server(servicer, server):
93115
rpc_method_handlers = {
@@ -111,6 +133,16 @@ def add_VectorDBServicer_to_server(servicer, server):
111133
request_deserializer=vector__db__pb2.SearchRequest.FromString,
112134
response_serializer=vector__db__pb2.SearchResponse.SerializeToString,
113135
),
136+
'InsertVectorsBatch': grpc.unary_unary_rpc_method_handler(
137+
servicer.InsertVectorsBatch,
138+
request_deserializer=vector__db__pb2.InsertVectorsBatchRequest.FromString,
139+
response_serializer=vector__db__pb2.InsertVectorsBatchResponse.SerializeToString,
140+
),
141+
'SearchPointsBatch': grpc.unary_unary_rpc_method_handler(
142+
servicer.SearchPointsBatch,
143+
request_deserializer=vector__db__pb2.SearchPointsBatchRequest.FromString,
144+
response_serializer=vector__db__pb2.SearchPointsBatchResponse.SerializeToString,
145+
),
114146
}
115147
generic_handler = grpc.method_handlers_generic_handler(
116148
'vectordb.VectorDB', rpc_method_handlers)
@@ -119,7 +151,7 @@ def add_VectorDBServicer_to_server(servicer, server):
119151

120152

121153
# This class is part of an EXPERIMENTAL API.
122-
class VectorDB(object):
154+
class VectorDB:
123155
"""Missing associated documentation comment in .proto file."""
124156

125157
@staticmethod
@@ -229,3 +261,57 @@ def SearchPoints(request,
229261
timeout,
230262
metadata,
231263
_registered_method=True)
264+
265+
@staticmethod
266+
def InsertVectorsBatch(request,
267+
target,
268+
options=(),
269+
channel_credentials=None,
270+
call_credentials=None,
271+
insecure=False,
272+
compression=None,
273+
wait_for_ready=None,
274+
timeout=None,
275+
metadata=None):
276+
return grpc.experimental.unary_unary(
277+
request,
278+
target,
279+
'/vectordb.VectorDB/InsertVectorsBatch',
280+
vector__db__pb2.InsertVectorsBatchRequest.SerializeToString,
281+
vector__db__pb2.InsertVectorsBatchResponse.FromString,
282+
options,
283+
channel_credentials,
284+
insecure,
285+
call_credentials,
286+
compression,
287+
wait_for_ready,
288+
timeout,
289+
metadata,
290+
_registered_method=True)
291+
292+
@staticmethod
293+
def SearchPointsBatch(request,
294+
target,
295+
options=(),
296+
channel_credentials=None,
297+
call_credentials=None,
298+
insecure=False,
299+
compression=None,
300+
wait_for_ready=None,
301+
timeout=None,
302+
metadata=None):
303+
return grpc.experimental.unary_unary(
304+
request,
305+
target,
306+
'/vectordb.VectorDB/SearchPointsBatch',
307+
vector__db__pb2.SearchPointsBatchRequest.SerializeToString,
308+
vector__db__pb2.SearchPointsBatchResponse.FromString,
309+
options,
310+
channel_credentials,
311+
insecure,
312+
call_credentials,
313+
compression,
314+
wait_for_ready,
315+
timeout,
316+
metadata,
317+
_registered_method=True)

0 commit comments

Comments
 (0)