From 5748ef9e696cec1e47945e3157ff507949c0a146 Mon Sep 17 00:00:00 2001 From: Arshdeep54 Date: Mon, 19 Oct 2026 05:08:51 +0530 Subject: [PATCH] add distance scores to dense search results --- .github/workflows/rust.yml | 2 +- client/python/USAGE.md | 6 +-- client/python/tests/test_async_client.py | 30 ++++++------ client/python/tests/test_client.py | 25 +++++----- client/python/vortexdb/async_client.py | 18 +++++--- client/python/vortexdb/client.py | 18 +++++--- client/python/vortexdb/grpc/vector_db_pb2.py | 48 ++++++++++---------- crates/api/src/lib.rs | 5 +- crates/api/src/tests.rs | 2 +- crates/defs/src/types.rs | 3 +- crates/grpc/proto/vector-db.proto | 1 + crates/grpc/src/service.rs | 41 +++++++++++------ crates/grpc/src/tests.rs | 1 + crates/http/src/handler.rs | 8 +++- crates/http/src/tests.rs | 1 + crates/index/src/flat/index.rs | 4 +- crates/index/src/flat/tests.rs | 35 ++++++++++++-- crates/index/src/hnsw/index.rs | 28 ++++++++---- crates/index/src/hnsw/tests.rs | 13 ++++-- crates/index/src/kd_tree/index.rs | 5 +- crates/index/src/kd_tree/tests.rs | 43 ++++++++++-------- crates/index/src/lib.rs | 7 ++- crates/tui/src/app/events.rs | 2 +- 23 files changed, 216 insertions(+), 130 deletions(-) diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index f0873be..e4e8810 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -20,7 +20,7 @@ jobs: - name: Set Up Rust uses: actions-rs/toolchain@v1 with: - toolchain: stable + toolchain: 1.98.1 components: clippy, rustfmt override: true diff --git a/client/python/USAGE.md b/client/python/USAGE.md index fbe22a0..18fd497 100644 --- a/client/python/USAGE.md +++ b/client/python/USAGE.md @@ -161,11 +161,11 @@ search( vector: DenseVector, similarity: Similarity, limit: int, -) -> list[str] +) -> list[tuple[str, float]] ``` Returns -- List of `point_id` strings +- List of `(point_id, score)` tuples, where score is the distance (lower is closer) Raises - `TypeError` if `vector` is not a `DenseVector` @@ -182,7 +182,7 @@ batch_search( queries, similarity: Similarity | None = None, limit: int | None = None, -) -> list[list[str]] +) -> list[list[tuple[str, float]]] ``` Returns diff --git a/client/python/tests/test_async_client.py b/client/python/tests/test_async_client.py index d2af9fd..dfbee27 100644 --- a/client/python/tests/test_async_client.py +++ b/client/python/tests/test_async_client.py @@ -173,7 +173,8 @@ async def run(): result_point_ids=[ Mock(id=Mock(value="p1")), Mock(id=Mock(value="p2")), - ] + ], + scores=[0.1, 0.2], ) results = await client.search( @@ -182,7 +183,7 @@ async def run(): limit=2, ) - assert results == ["p1", "p2"] + assert results == [("p1", 0.1), ("p2", 0.2)] asyncio.run(run()) @@ -192,20 +193,21 @@ async def run(): mock_connection.call.return_value = Mock( result_point_ids=[ Mock(id=Mock(value="p1")), - ] + ], + scores=[0.1], ) q = SearchQuery(DenseVector([1, 2, 3]), Similarity.COSINE, 2) results = await client.search(query=q) - assert results == ["p1"] + assert results == [("p1", 0.1)] asyncio.run(run()) def test_async_search_accepts_ef(client, mock_connection): async def run(): - mock_connection.call.return_value = Mock(result_point_ids=[]) + mock_connection.call.return_value = Mock(result_point_ids=[], scores=[]) await client.search( vector=DenseVector([1, 2, 3]), @@ -239,8 +241,8 @@ def test_async_batch_search_full_tuple(client, mock_connection): async def run(): mock_connection.call.return_value = Mock( results=[ - Mock(result_point_ids=[Mock(id=Mock(value="p1"))]), - Mock(result_point_ids=[Mock(id=Mock(value="p2"))]), + Mock(result_point_ids=[Mock(id=Mock(value="p1"))], scores=[0.5]), + Mock(result_point_ids=[Mock(id=Mock(value="p2"))], scores=[0.5]), ] ) queries = [ @@ -248,7 +250,7 @@ async def run(): (DenseVector([4, 5, 6]), Similarity.EUCLIDEAN, 1), ] result = await client.batch_search(queries=queries) - assert result == [["p1"], ["p2"]] + assert result == [[("p1", 0.5)], [("p2", 0.5)]] asyncio.run(run()) @@ -257,14 +259,14 @@ def test_async_batch_search_searchquery_objects(client, mock_connection): async def run(): mock_connection.call.return_value = Mock( results=[ - Mock(result_point_ids=[Mock(id=Mock(value="p1"))]), + Mock(result_point_ids=[Mock(id=Mock(value="p1"))], scores=[0.5]), ] ) queries = [ SearchQuery(DenseVector([1, 2, 3]), Similarity.COSINE, 2), ] result = await client.batch_search(queries=queries) - assert result == [["p1"]] + assert result == [[("p1", 0.5)]] asyncio.run(run()) @@ -273,7 +275,7 @@ def test_async_batch_search_vectors_with_global_params(client, mock_connection): async def run(): mock_connection.call.return_value = Mock( results=[ - Mock(result_point_ids=[Mock(id=Mock(value="p1"))]), + Mock(result_point_ids=[Mock(id=Mock(value="p1"))], scores=[0.5]), ] ) queries = [DenseVector([1, 2, 3])] @@ -282,7 +284,7 @@ async def run(): similarity=Similarity.MANHATTAN, limit=2, ) - assert result == [["p1"]] + assert result == [[("p1", 0.5)]] asyncio.run(run()) @@ -293,7 +295,7 @@ def test_async_batch_search_vector_similarity_with_global_limit( async def run(): mock_connection.call.return_value = Mock( results=[ - Mock(result_point_ids=[Mock(id=Mock(value="p1"))]), + Mock(result_point_ids=[Mock(id=Mock(value="p1"))], scores=[0.5]), ] ) queries = [ @@ -303,7 +305,7 @@ async def run(): queries=queries, limit=2, ) - assert result == [["p1"]] + assert result == [[("p1", 0.5)]] asyncio.run(run()) diff --git a/client/python/tests/test_client.py b/client/python/tests/test_client.py index b639da5..5fe3518 100644 --- a/client/python/tests/test_client.py +++ b/client/python/tests/test_client.py @@ -139,7 +139,8 @@ def test_search_success(client, mock_connection): result_point_ids=[ Mock(id=Mock(value="p1")), Mock(id=Mock(value="p2")), - ] + ], + scores=[0.1, 0.2], ) results = client.search( @@ -148,11 +149,11 @@ def test_search_success(client, mock_connection): limit=2, ) - assert results == ["p1", "p2"] + assert results == [("p1", 0.1), ("p2", 0.2)] def test_search_accepts_ef(client, mock_connection): - mock_connection.call.return_value = Mock(result_point_ids=[]) + mock_connection.call.return_value = Mock(result_point_ids=[], scores=[]) client.search( vector=DenseVector([1, 2, 3]), @@ -195,8 +196,8 @@ def test_batch_search_accepts_ef(client, mock_connection): def test_batch_search_full_tuple(client, mock_connection): mock_connection.call.return_value = Mock( results=[ - Mock(result_point_ids=[Mock(id=Mock(value="p1"))]), - Mock(result_point_ids=[Mock(id=Mock(value="p2"))]), + Mock(result_point_ids=[Mock(id=Mock(value="p1"))], scores=[0.5]), + Mock(result_point_ids=[Mock(id=Mock(value="p2"))], scores=[0.5]), ] ) queries = [ @@ -204,13 +205,13 @@ def test_batch_search_full_tuple(client, mock_connection): (DenseVector([4, 5, 6]), Similarity.EUCLIDEAN, 1), ] result = client.batch_search(queries=queries) - assert result == [["p1"], ["p2"]] + assert result == [[("p1", 0.5)], [("p2", 0.5)]] def test_batch_search_vectors_with_global_params(client, mock_connection): mock_connection.call.return_value = Mock( results=[ - Mock(result_point_ids=[Mock(id=Mock(value="p1"))]), + Mock(result_point_ids=[Mock(id=Mock(value="p1"))], scores=[0.5]), ] ) queries = [DenseVector([1, 2, 3])] @@ -219,13 +220,13 @@ def test_batch_search_vectors_with_global_params(client, mock_connection): similarity=Similarity.MANHATTAN, limit=2, ) - assert result == [["p1"]] + assert result == [[("p1", 0.5)]] def test_batch_search_vector_similarity_with_global_limit(client, mock_connection): mock_connection.call.return_value = Mock( results=[ - Mock(result_point_ids=[Mock(id=Mock(value="p1"))]), + Mock(result_point_ids=[Mock(id=Mock(value="p1"))], scores=[0.5]), ] ) queries = [ @@ -235,20 +236,20 @@ def test_batch_search_vector_similarity_with_global_limit(client, mock_connectio queries=queries, limit=2, ) - assert result == [["p1"]] + assert result == [[("p1", 0.5)]] def test_batch_search_searchquery_objects(client, mock_connection): mock_connection.call.return_value = Mock( results=[ - Mock(result_point_ids=[Mock(id=Mock(value="p1"))]), + Mock(result_point_ids=[Mock(id=Mock(value="p1"))], scores=[0.5]), ] ) queries = [ SearchQuery(DenseVector([1, 2, 3]), Similarity.COSINE, 2), ] result = client.batch_search(queries=queries) - assert result == [["p1"]] + assert result == [[("p1", 0.5)]] def test_batch_search_missing_globals_for_vector(client): diff --git a/client/python/vortexdb/async_client.py b/client/python/vortexdb/async_client.py index 3b7c2f5..3ce0aaf 100644 --- a/client/python/vortexdb/async_client.py +++ b/client/python/vortexdb/async_client.py @@ -1,4 +1,4 @@ -from typing import List +from typing import List, Tuple from vortexdb import protoutils as proto from vortexdb.async_connection import AsyncGRPCConnection @@ -98,10 +98,10 @@ async def search( limit: int | None = None, query: SearchQuery | None = None, ef: int | None = None, - ) -> List[str]: + ) -> List[Tuple[str, float]]: """ Search for nearest neighbors. - Returns: List of point IDs + Returns: List of (point ID, score) tuples """ if query is not None: if not isinstance(query, SearchQuery): @@ -131,7 +131,10 @@ async def search( request, ) - return [pid.id.value for pid in response.result_point_ids] + return [ + (pid.id.value, score) + for pid, score in zip(response.result_point_ids, response.scores) + ] async def batch_search( self, @@ -140,7 +143,7 @@ async def batch_search( similarity: Similarity | None = None, limit: int | None = None, ef: int | None = None, - ) -> List[List[str]]: + ) -> List[List[Tuple[str, float]]]: """ Flexible batch search. @@ -194,7 +197,10 @@ async def batch_search( request = proto.build_batch_search_request(queries=normalized, ef=ef) response = await self._conn.call(self._conn.stub.SearchPointsBatch, request) return [ - [pid.id.value for pid in result.result_point_ids] + [ + (pid.id.value, score) + for pid, score in zip(result.result_point_ids, result.scores) + ] for result in response.results ] diff --git a/client/python/vortexdb/client.py b/client/python/vortexdb/client.py index 9984279..f38c934 100644 --- a/client/python/vortexdb/client.py +++ b/client/python/vortexdb/client.py @@ -1,4 +1,4 @@ -from typing import List +from typing import List, Tuple from vortexdb.connection import GRPCConnection from vortexdb.config import VortexDBConfig @@ -101,10 +101,10 @@ def search( limit: int | None = None, query: SearchQuery | None = None, ef: int | None = None, - ) -> List[str]: + ) -> List[Tuple[str, float]]: """ Search for nearest neighbors. - Returns: List of point IDs + Returns: List of (point ID, score) tuples """ if query is not None: if not isinstance(query, SearchQuery): @@ -129,7 +129,10 @@ def search( self._conn.stub.SearchPoints, request, ) - return [pid.id.value for pid in response.result_point_ids] + return [ + (pid.id.value, score) + for pid, score in zip(response.result_point_ids, response.scores) + ] def batch_search( self, @@ -138,7 +141,7 @@ def batch_search( similarity: Similarity | None = None, limit: int | None = None, ef: int | None = None, - ) -> List[List[str]]: + ) -> List[List[Tuple[str, float]]]: """ Flexible batch search. @@ -192,7 +195,10 @@ def batch_search( request = proto.build_batch_search_request(queries=normalized, ef=ef) response = self._conn.call(self._conn.stub.SearchPointsBatch, request) return [ - [pid.id.value for pid in result.result_point_ids] + [ + (pid.id.value, score) + for pid, score in zip(result.result_point_ids, result.scores) + ] for result in response.results ] diff --git a/client/python/vortexdb/grpc/vector_db_pb2.py b/client/python/vortexdb/grpc/vector_db_pb2.py index faa4212..3cdf444 100644 --- a/client/python/vortexdb/grpc/vector_db_pb2.py +++ b/client/python/vortexdb/grpc/vector_db_pb2.py @@ -25,17 +25,17 @@ from google.protobuf import empty_pb2 as google_dot_protobuf_dot_empty__pb2 -DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x0fvector-db.proto\x12\x08vectordb\x1a\x1bgoogle/protobuf/empty.proto\"\x15\n\x04UUID\x12\r\n\x05value\x18\x01 \x01(\t\"`\n\x13InsertVectorRequest\x12%\n\x06vector\x18\x01 \x01(\x0b\x32\x15.vectordb.DenseVector\x12\"\n\x07payload\x18\x02 \x01(\x0b\x32\x11.vectordb.Payload\"\x81\x01\n\rSearchRequest\x12+\n\x0cquery_vector\x18\x01 \x01(\x0b\x32\x15.vectordb.DenseVector\x12(\n\nsimilarity\x18\x02 \x01(\x0e\x32\x14.vectordb.Similarity\x12\r\n\x05limit\x18\x03 \x01(\x04\x12\n\n\x02\x65\x66\x18\x04 \x01(\x04\"=\n\x0eSearchResponse\x12+\n\x10result_point_ids\x18\x01 \x03(\x0b\x32\x11.vectordb.PointID\"\x1d\n\x0b\x44\x65nseVector\x12\x0e\n\x06values\x18\x01 \x03(\x02\"q\n\x05Point\x12\x1d\n\x02id\x18\x01 \x01(\x0b\x32\x11.vectordb.PointID\x12\"\n\x07payload\x18\x02 \x01(\x0b\x32\x11.vectordb.Payload\x12%\n\x06vector\x18\x03 \x01(\x0b\x32\x15.vectordb.DenseVector\"%\n\x07PointID\x12\x1a\n\x02id\x18\x01 \x01(\x0b\x32\x0e.vectordb.UUID\"G\n\x07Payload\x12+\n\x0c\x63ontent_type\x18\x01 \x01(\x0e\x32\x15.vectordb.ContentType\x12\x0f\n\x07\x63ontent\x18\x02 \x01(\t\"K\n\x19InsertVectorsBatchRequest\x12.\n\x07vectors\x18\x01 \x03(\x0b\x32\x1d.vectordb.InsertVectorRequest\"<\n\x1aInsertVectorsBatchResponse\x12\x1e\n\x03ids\x18\x01 \x03(\x0b\x32\x11.vectordb.PointID\"D\n\x18SearchPointsBatchRequest\x12(\n\x07queries\x18\x01 \x03(\x0b\x32\x17.vectordb.SearchRequest\"F\n\x19SearchPointsBatchResponse\x12)\n\x07results\x18\x01 \x03(\x0b\x32\x18.vectordb.SearchResponse*C\n\nSimilarity\x12\r\n\tEuclidean\x10\x00\x12\r\n\tManhattan\x10\x01\x12\x0b\n\x07Hamming\x10\x02\x12\n\n\x06\x43osine\x10\x03*\"\n\x0b\x43ontentType\x12\t\n\x05Image\x10\x00\x12\x08\n\x04Text\x10\x01\x32\xc4\x03\n\x08VectorDB\x12\x42\n\x0cInsertVector\x12\x1d.vectordb.InsertVectorRequest\x1a\x11.vectordb.PointID\"\x00\x12:\n\x0b\x44\x65letePoint\x12\x11.vectordb.PointID\x1a\x16.google.protobuf.Empty\"\x00\x12\x30\n\x08GetPoint\x12\x11.vectordb.PointID\x1a\x0f.vectordb.Point\"\x00\x12\x43\n\x0cSearchPoints\x12\x17.vectordb.SearchRequest\x1a\x18.vectordb.SearchResponse\"\x00\x12\x61\n\x12InsertVectorsBatch\x12#.vectordb.InsertVectorsBatchRequest\x1a$.vectordb.InsertVectorsBatchResponse\"\x00\x12^\n\x11SearchPointsBatch\x12\".vectordb.SearchPointsBatchRequest\x1a#.vectordb.SearchPointsBatchResponse\"\x00\x62\x06proto3') +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x0fvector-db.proto\x12\x08vectordb\x1a\x1bgoogle/protobuf/empty.proto\"\x15\n\x04UUID\x12\r\n\x05value\x18\x01 \x01(\t\"`\n\x13InsertVectorRequest\x12%\n\x06vector\x18\x01 \x01(\x0b\x32\x15.vectordb.DenseVector\x12\"\n\x07payload\x18\x02 \x01(\x0b\x32\x11.vectordb.Payload\"\x81\x01\n\rSearchRequest\x12+\n\x0cquery_vector\x18\x01 \x01(\x0b\x32\x15.vectordb.DenseVector\x12(\n\nsimilarity\x18\x02 \x01(\x0e\x32\x14.vectordb.Similarity\x12\r\n\x05limit\x18\x03 \x01(\x04\x12\n\n\x02\x65\x66\x18\x04 \x01(\x04\"M\n\x0eSearchResponse\x12+\n\x10result_point_ids\x18\x01 \x03(\x0b\x32\x11.vectordb.PointID\x12\x0e\n\x06scores\x18\x02 \x03(\x02\"\x1d\n\x0b\x44\x65nseVector\x12\x0e\n\x06values\x18\x01 \x03(\x02\"q\n\x05Point\x12\x1d\n\x02id\x18\x01 \x01(\x0b\x32\x11.vectordb.PointID\x12\"\n\x07payload\x18\x02 \x01(\x0b\x32\x11.vectordb.Payload\x12%\n\x06vector\x18\x03 \x01(\x0b\x32\x15.vectordb.DenseVector\"%\n\x07PointID\x12\x1a\n\x02id\x18\x01 \x01(\x0b\x32\x0e.vectordb.UUID\"G\n\x07Payload\x12+\n\x0c\x63ontent_type\x18\x01 \x01(\x0e\x32\x15.vectordb.ContentType\x12\x0f\n\x07\x63ontent\x18\x02 \x01(\t\"K\n\x19InsertVectorsBatchRequest\x12.\n\x07vectors\x18\x01 \x03(\x0b\x32\x1d.vectordb.InsertVectorRequest\"<\n\x1aInsertVectorsBatchResponse\x12\x1e\n\x03ids\x18\x01 \x03(\x0b\x32\x11.vectordb.PointID\"D\n\x18SearchPointsBatchRequest\x12(\n\x07queries\x18\x01 \x03(\x0b\x32\x17.vectordb.SearchRequest\"F\n\x19SearchPointsBatchResponse\x12)\n\x07results\x18\x01 \x03(\x0b\x32\x18.vectordb.SearchResponse*C\n\nSimilarity\x12\r\n\tEuclidean\x10\x00\x12\r\n\tManhattan\x10\x01\x12\x0b\n\x07Hamming\x10\x02\x12\n\n\x06\x43osine\x10\x03*\"\n\x0b\x43ontentType\x12\t\n\x05Image\x10\x00\x12\x08\n\x04Text\x10\x01\x32\xc4\x03\n\x08VectorDB\x12\x42\n\x0cInsertVector\x12\x1d.vectordb.InsertVectorRequest\x1a\x11.vectordb.PointID\"\x00\x12:\n\x0b\x44\x65letePoint\x12\x11.vectordb.PointID\x1a\x16.google.protobuf.Empty\"\x00\x12\x30\n\x08GetPoint\x12\x11.vectordb.PointID\x1a\x0f.vectordb.Point\"\x00\x12\x43\n\x0cSearchPoints\x12\x17.vectordb.SearchRequest\x1a\x18.vectordb.SearchResponse\"\x00\x12\x61\n\x12InsertVectorsBatch\x12#.vectordb.InsertVectorsBatchRequest\x1a$.vectordb.InsertVectorsBatchResponse\"\x00\x12^\n\x11SearchPointsBatch\x12\".vectordb.SearchPointsBatchRequest\x1a#.vectordb.SearchPointsBatchResponse\"\x00\x62\x06proto3') _globals = globals() _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) _builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'vector_db_pb2', _globals) if not _descriptor._USE_C_DESCRIPTORS: DESCRIPTOR._loaded_options = None - _globals['_SIMILARITY']._serialized_start=913 - _globals['_SIMILARITY']._serialized_end=980 - _globals['_CONTENTTYPE']._serialized_start=982 - _globals['_CONTENTTYPE']._serialized_end=1016 + _globals['_SIMILARITY']._serialized_start=929 + _globals['_SIMILARITY']._serialized_end=996 + _globals['_CONTENTTYPE']._serialized_start=998 + _globals['_CONTENTTYPE']._serialized_end=1032 _globals['_UUID']._serialized_start=58 _globals['_UUID']._serialized_end=79 _globals['_INSERTVECTORREQUEST']._serialized_start=81 @@ -43,23 +43,23 @@ _globals['_SEARCHREQUEST']._serialized_start=180 _globals['_SEARCHREQUEST']._serialized_end=309 _globals['_SEARCHRESPONSE']._serialized_start=311 - _globals['_SEARCHRESPONSE']._serialized_end=372 - _globals['_DENSEVECTOR']._serialized_start=374 - _globals['_DENSEVECTOR']._serialized_end=403 - _globals['_POINT']._serialized_start=405 - _globals['_POINT']._serialized_end=518 - _globals['_POINTID']._serialized_start=520 - _globals['_POINTID']._serialized_end=557 - _globals['_PAYLOAD']._serialized_start=559 - _globals['_PAYLOAD']._serialized_end=630 - _globals['_INSERTVECTORSBATCHREQUEST']._serialized_start=632 - _globals['_INSERTVECTORSBATCHREQUEST']._serialized_end=707 - _globals['_INSERTVECTORSBATCHRESPONSE']._serialized_start=709 - _globals['_INSERTVECTORSBATCHRESPONSE']._serialized_end=769 - _globals['_SEARCHPOINTSBATCHREQUEST']._serialized_start=771 - _globals['_SEARCHPOINTSBATCHREQUEST']._serialized_end=839 - _globals['_SEARCHPOINTSBATCHRESPONSE']._serialized_start=841 - _globals['_SEARCHPOINTSBATCHRESPONSE']._serialized_end=911 - _globals['_VECTORDB']._serialized_start=1019 - _globals['_VECTORDB']._serialized_end=1471 + _globals['_SEARCHRESPONSE']._serialized_end=388 + _globals['_DENSEVECTOR']._serialized_start=390 + _globals['_DENSEVECTOR']._serialized_end=419 + _globals['_POINT']._serialized_start=421 + _globals['_POINT']._serialized_end=534 + _globals['_POINTID']._serialized_start=536 + _globals['_POINTID']._serialized_end=573 + _globals['_PAYLOAD']._serialized_start=575 + _globals['_PAYLOAD']._serialized_end=646 + _globals['_INSERTVECTORSBATCHREQUEST']._serialized_start=648 + _globals['_INSERTVECTORSBATCHREQUEST']._serialized_end=723 + _globals['_INSERTVECTORSBATCHRESPONSE']._serialized_start=725 + _globals['_INSERTVECTORSBATCHRESPONSE']._serialized_end=785 + _globals['_SEARCHPOINTSBATCHREQUEST']._serialized_start=787 + _globals['_SEARCHPOINTSBATCHREQUEST']._serialized_end=855 + _globals['_SEARCHPOINTSBATCHRESPONSE']._serialized_start=857 + _globals['_SEARCHPOINTSBATCHRESPONSE']._serialized_end=927 + _globals['_VECTORDB']._serialized_start=1035 + _globals['_VECTORDB']._serialized_end=1487 # @@protoc_insertion_point(module_scope) diff --git a/crates/api/src/lib.rs b/crates/api/src/lib.rs index 741efb2..5377dde 100644 --- a/crates/api/src/lib.rs +++ b/crates/api/src/lib.rs @@ -125,7 +125,7 @@ impl VectorDb { } } - pub fn search(&self, query: SearchQueryInput) -> Result> { + pub fn search(&self, query: SearchQueryInput) -> Result> { // Validate search limit if query.limit == 0 { return Err(ApiError::InvalidSearchLimit { limit: query.limit }); @@ -142,14 +142,13 @@ impl VectorDb { // Use vector index to find similar vectors let index = self.index.read().map_err(|_| ApiError::LockError)?; - //TODO: Add feat of returning similarity scores in the search let vectors = index.search_with_ef(query.vector, query.similarity, query.limit, query.ef)?; Ok(vectors) } - pub fn search_batch(&self, queries: Vec) -> Result>> { + pub fn search_batch(&self, queries: Vec) -> Result>> { let mut results = Vec::with_capacity(queries.len()); let index = self.index.read().unwrap(); diff --git a/crates/api/src/tests.rs b/crates/api/src/tests.rs index ea9504b..35df694 100644 --- a/crates/api/src/tests.rs +++ b/crates/api/src/tests.rs @@ -154,7 +154,7 @@ fn test_search() { .unwrap(); assert_eq!(results.len(), 1); - assert_eq!(results[0], ids[0]); // The first vector should be closest + assert_eq!(results[0].0, ids[0]); // The first vector should be closest } #[test] diff --git a/crates/defs/src/types.rs b/crates/defs/src/types.rs index 1507653..7e67399 100644 --- a/crates/defs/src/types.rs +++ b/crates/defs/src/types.rs @@ -55,7 +55,7 @@ pub struct IndexedVector { pub vector: DenseVector, } -#[derive(Debug, Serialize, Deserialize, Copy, Clone)] +#[derive(Debug, Serialize, Deserialize, Copy, Clone, PartialEq, Eq)] pub enum Similarity { Euclidean, Manhattan, @@ -92,6 +92,7 @@ pub struct SearchQueryInput { #[derive(Clone, Serialize, Deserialize, Debug)] pub struct SearchResponse { pub results: Vec, + pub scores: Vec, } #[derive(Debug, Clone, Serialize, Deserialize)] diff --git a/crates/grpc/proto/vector-db.proto b/crates/grpc/proto/vector-db.proto index 0b5dfe1..ecca463 100644 --- a/crates/grpc/proto/vector-db.proto +++ b/crates/grpc/proto/vector-db.proto @@ -42,6 +42,7 @@ message SearchRequest { message SearchResponse { repeated PointID result_point_ids = 1; + repeated float scores = 2; } message DenseVector { diff --git a/crates/grpc/src/service.rs b/crates/grpc/src/service.rs index 9225077..9807bf2 100644 --- a/crates/grpc/src/service.rs +++ b/crates/grpc/src/service.rs @@ -142,17 +142,23 @@ impl VectorDb for VectorDBService { .map_err(|e| Status::from(crate::error::GrpcError::from(e)))?; // create a mapped vector of PointIds - let result = result_point_ids + let (result, scores) = result_point_ids .into_iter() - .map(|id| PointId { - id: Some(Uuid { - value: id.to_string(), - }), + .map(|(id, score)| { + ( + PointId { + id: Some(Uuid { + value: id.to_string(), + }), + }, + score, + ) }) - .collect(); + .unzip(); Ok(Response::new(SearchResponse { result_point_ids: result, + scores, })) } @@ -239,15 +245,22 @@ impl VectorDb for VectorDBService { }) .map_err(|e| tonic::Status::from(GrpcError::from(e)))?; + let (result_point_ids, scores) = ids + .into_iter() + .map(|(id, score)| { + ( + PointId { + id: Some(Uuid { + value: id.to_string(), + }), + }, + score, + ) + }) + .unzip(); results.push(SearchResponse { - result_point_ids: ids - .into_iter() - .map(|id| PointId { - id: Some(Uuid { - value: id.to_string(), - }), - }) - .collect(), + result_point_ids, + scores, }); } diff --git a/crates/grpc/src/tests.rs b/crates/grpc/src/tests.rs index e623e13..006649b 100644 --- a/crates/grpc/src/tests.rs +++ b/crates/grpc/src/tests.rs @@ -242,6 +242,7 @@ async fn test_search_vector_rpc() { // check if the returned point id matches the inserted point id assert_eq!(result.result_point_ids[0], PointId { id: point.id }); + assert_eq!(result.scores, vec![2f32.sqrt()]); } #[tokio::test] diff --git a/crates/http/src/handler.rs b/crates/http/src/handler.rs index 387d167..e8e171d 100644 --- a/crates/http/src/handler.rs +++ b/crates/http/src/handler.rs @@ -98,7 +98,8 @@ pub async fn search_points_handler( ) -> Result, (StatusCode, String)> { match app_state.db.search(request) { Ok(results) => { - let response = SearchResponse { results }; + let (results, scores) = results.into_iter().unzip(); + let response = SearchResponse { results, scores }; Ok(Json(response)) } Err(e) => { @@ -117,7 +118,10 @@ pub async fn batch_search_handler( .search_batch(request.queries) .map_err(|e| AppError::Api(e.to_string()))? .into_iter() - .map(|ids| SearchResponse { results: ids }) + .map(|found| { + let (results, scores) = found.into_iter().unzip(); + SearchResponse { results, scores } + }) .collect(); Ok(Json(BatchSearchResponse { results })) diff --git a/crates/http/src/tests.rs b/crates/http/src/tests.rs index b201716..204ed8b 100644 --- a/crates/http/src/tests.rs +++ b/crates/http/src/tests.rs @@ -84,6 +84,7 @@ async fn in_memory_storage_http_smoke_test() { search_response.assert_status_ok(); let search_body: serde_json::Value = search_response.json(); assert_eq!(search_body["results"], json!([point_id])); + assert_eq!(search_body["scores"], json!([0.0])); let delete_response = server .delete(&format!("/points/{point_id}")) diff --git a/crates/index/src/flat/index.rs b/crates/index/src/flat/index.rs index 927e41c..9758756 100644 --- a/crates/index/src/flat/index.rs +++ b/crates/index/src/flat/index.rs @@ -41,7 +41,7 @@ impl VectorIndex for FlatIndex { query_vector: DenseVector, similarity: Similarity, k: usize, - ) -> Result, IndexError> { + ) -> Result, IndexError> { let scores = self .index .iter() @@ -65,7 +65,7 @@ impl VectorIndex for FlatIndex { Ok(heap .into_sorted_vec() .into_iter() - .map(|v| v.point_id.unwrap()) + .map(|v| (v.point_id.unwrap(), v.distance)) .collect()) } } diff --git a/crates/index/src/flat/tests.rs b/crates/index/src/flat/tests.rs index 6d43c3d..6f523e6 100644 --- a/crates/index/src/flat/tests.rs +++ b/crates/index/src/flat/tests.rs @@ -95,7 +95,7 @@ fn test_search_euclidean() { let results = index .search(vec![0.0, 0.0], Similarity::Euclidean, 2) .unwrap(); - assert_eq!(results, vec![id1, id2]); + assert_eq!(results, vec![(id1, 2f32.sqrt()), (id2, 8f32.sqrt())]); } #[test] @@ -124,7 +124,33 @@ fn test_search_cosine() { .unwrap(); let results = index.search(vec![1.0, 1.0], Similarity::Cosine, 2).unwrap(); - assert_eq!(results, vec![id2, id1]); + let ids: Vec<_> = results.into_iter().map(|(id, _)| id).collect(); + assert_eq!(ids, vec![id2, id1]); +} + +#[test] +fn test_search_cosine_zero_vector() { + let mut index = FlatIndex::new(); + let id1 = Uuid::new_v4(); + let id2 = Uuid::new_v4(); + index + .insert(IndexedVector { + id: id1, + vector: vec![1.0, 0.0], + }) + .unwrap(); + index + .insert(IndexedVector { + id: id2, + vector: vec![0.0, 0.0], + }) + .unwrap(); + + let results = index.search(vec![1.0, 0.0], Similarity::Cosine, 2).unwrap(); + assert_eq!(results, vec![(id1, 0.0), (id2, 1.0)]); + + let results = index.search(vec![0.0, 0.0], Similarity::Cosine, 2).unwrap(); + assert!(results.iter().all(|(_, d)| *d == 1.0)); } #[test] @@ -155,7 +181,7 @@ fn test_search_manhattan() { let results = index .search(vec![0.0, 0.0], Similarity::Manhattan, 2) .unwrap(); - assert_eq!(results, vec![id1, id2]); + assert_eq!(results, vec![(id1, 2.0), (id2, 4.0)]); } #[test] @@ -186,7 +212,8 @@ fn test_search_hamming() { let results = index .search(vec![1.0, 0.0, 0.0, 0.0], Similarity::Hamming, 2) .unwrap(); - assert_eq!(results, vec![id2, id3]); + let ids: Vec<_> = results.into_iter().map(|(id, _)| id).collect(); + assert_eq!(ids, vec![id2, id3]); } #[test] diff --git a/crates/index/src/hnsw/index.rs b/crates/index/src/hnsw/index.rs index f730198..1c1646c 100644 --- a/crates/index/src/hnsw/index.rs +++ b/crates/index/src/hnsw/index.rs @@ -1,7 +1,6 @@ use std::collections::HashMap; use defs::{DenseVector, Dimension, IndexedVector, PointId, Similarity}; -use uuid::Uuid; use crate::VectorIndex; use crate::{IndexError, Result, distance}; @@ -197,23 +196,34 @@ impl VectorIndex for HnswIndex { } Ok(false) } - /// Search for top-k ids + /// Search for top-k (id, distance) pairs /// - normalize query for cosine (1 − cos) /// - pick a non-deleted entry point /// - greedy descend from the top layer to level 1 /// - run ef-best-first at level 0 with ef0 = max(ef, k) - /// - return up to k ids by ascending distance - fn search(&self, query: DenseVector, similarity: Similarity, k: usize) -> Result> { + /// - return up to k (id, distance) pairs by ascending distance + fn search( + &self, + query: DenseVector, + similarity: Similarity, + k: usize, + ) -> Result> { self.search_with_ef(query, similarity, k, None) } fn search_with_ef( &self, mut query: DenseVector, - _similarity: Similarity, + similarity: Similarity, k: usize, ef: Option, - ) -> Result> { + ) -> Result> { + if similarity != self.similarity { + return Err(IndexError::UnsupportedSimilarity { + metric: format!("{:?}", similarity), + }); + } + if k == 0 { return Ok(Vec::new()); } @@ -247,8 +257,10 @@ impl VectorIndex for HnswIndex { let ef0 = max(ef.unwrap_or(self.ef), k); let mut w = self.search_layer_for_insert(ep, 0, &query, ef0)?; w.truncate(k); - let result: Vec = w.into_iter().map(|(id, _)| id).collect(); - Ok(result) + if matches!(self.similarity, Similarity::Euclidean) { + w.iter_mut().for_each(|(_, d)| *d = d.sqrt()); + } + Ok(w) } } diff --git a/crates/index/src/hnsw/tests.rs b/crates/index/src/hnsw/tests.rs index 605c8f7..76f942d 100644 --- a/crates/index/src/hnsw/tests.rs +++ b/crates/index/src/hnsw/tests.rs @@ -1,6 +1,6 @@ use super::*; -use crate::VectorIndex; use crate::flat::index::FlatIndex; +use crate::{IndexError, VectorIndex}; use defs::{IndexedVector, Similarity}; use uuid::Uuid; @@ -116,6 +116,13 @@ fn test_search_empty_index_returns_empty() { assert!(res.is_empty()); } +#[test] +fn test_search_rejects_mismatched_similarity() { + let index = HnswIndex::new(Similarity::Euclidean, 2); + let res = index.search(vec![0.0, 0.0], Similarity::Manhattan, 1); + assert!(matches!(res, Err(IndexError::UnsupportedSimilarity { .. }))); +} + #[test] fn test_search_k_zero_returns_empty() { let mut index = HnswIndex::new(Similarity::Euclidean, 2); @@ -149,7 +156,7 @@ fn test_search_cosine_normalization_basic() { let res = index .search(vec![10.0, 0.0], Similarity::Cosine, 1) .unwrap(); - assert_eq!(res, vec![ID1]); + assert_eq!(res, vec![(ID1, 0.0)]); } #[test] @@ -183,7 +190,7 @@ fn test_soft_delete_and_search_skip() { let res = index .search(vec![0.9, 0.1], Similarity::Euclidean, 2) .unwrap(); - assert!(!res.contains(&ID2)); + assert!(!res.iter().any(|(id, _)| *id == ID2)); // Deleting a non-existent id returns false let existed = index.delete(ID4).unwrap(); diff --git a/crates/index/src/kd_tree/index.rs b/crates/index/src/kd_tree/index.rs index 31b6c6e..65301a1 100644 --- a/crates/index/src/kd_tree/index.rs +++ b/crates/index/src/kd_tree/index.rs @@ -457,14 +457,13 @@ impl VectorIndex for KDTree { query_vector: DenseVector, similarity: Similarity, k: usize, - ) -> Result> { + ) -> Result> { if matches!(similarity, Similarity::Cosine | Similarity::Hamming) { return Err(IndexError::UnsupportedSimilarity { metric: format!("{:?}", similarity), }); } - let results = self.search_top_k(query_vector, k, similarity); - Ok(results.into_iter().map(|(id, _)| id).collect()) + Ok(self.search_top_k(query_vector, k, similarity)) } } diff --git a/crates/index/src/kd_tree/tests.rs b/crates/index/src/kd_tree/tests.rs index a5425c8..6ab45e0 100644 --- a/crates/index/src/kd_tree/tests.rs +++ b/crates/index/src/kd_tree/tests.rs @@ -164,8 +164,8 @@ fn test_deleted_point_not_in_search_results() { let results = tree .search(vec![0.0, 0.0], Similarity::Euclidean, 2) .unwrap(); - assert!(!results.contains(&id1)); - assert!(results.contains(&id2)); + assert!(!results.iter().any(|(r, _)| *r == id1)); + assert!(results.iter().any(|(r, _)| *r == id2)); } // Search Tests (VectorIndex trait) @@ -195,8 +195,8 @@ fn test_search_euclidean() { .search(vec![0.0, 0.0], Similarity::Euclidean, 2) .unwrap(); assert_eq!(results.len(), 2); - assert_eq!(results[0], id1); // Closest - assert_eq!(results[1], id2); // Second closest + assert_eq!(results[0].0, id1); // Closest + assert_eq!(results[1].0, id2); // Second closest } #[test] @@ -215,8 +215,8 @@ fn test_search_manhattan() { .search(vec![0.0, 0.0], Similarity::Manhattan, 2) .unwrap(); assert_eq!(results.len(), 2); - assert_eq!(results[0], id1); - assert_eq!(results[1], id2); + assert_eq!(results[0].0, id1); + assert_eq!(results[1].0, id2); } #[test] @@ -273,7 +273,7 @@ fn test_search_exact_match() { .search(vec![5.0, 5.0], Similarity::Euclidean, 1) .unwrap(); assert_eq!(results.len(), 1); - assert_eq!(results[0], id1); + assert_eq!(results[0].0, id1); } // Search Correctness Tests @@ -299,7 +299,7 @@ fn test_search_correctness_3d() { // id2 at distance sqrt(0.75) ≈ 0.866 // Both are equidistant, should return both assert_eq!(results.len(), 2); - assert!(results.contains(&id1) || results.contains(&id2)); + assert!(results.iter().any(|(r, _)| *r == id1 || *r == id2)); } #[test] @@ -319,8 +319,8 @@ fn test_search_after_insert() { let results = tree .search(vec![0.0, 0.0], Similarity::Euclidean, 2) .unwrap(); - assert_eq!(results[0], id2); // Closest to origin - assert_eq!(results[1], id3); // Second closest + assert_eq!(results[0].0, id2); // Closest to origin + assert_eq!(results[1].0, id3); // Second closest } #[test] @@ -337,7 +337,7 @@ fn test_search_high_dimensional() { let query = vec![0.1; dim]; let results = tree.search(query, Similarity::Euclidean, 1).unwrap(); - assert_eq!(results[0], id1); // Closer to all-zeros + assert_eq!(results[0].0, id1); // Closer to all-zeros } // Rebalancing Tests @@ -361,7 +361,7 @@ fn test_many_inserts_maintains_searchability() { .unwrap(); assert_eq!(results.len(), 5); // First result should be the point at (0, 0) - assert_eq!(results[0], ids[0]); + assert_eq!(results[0].0, ids[0]); } #[test] @@ -389,7 +389,7 @@ fn test_delete_triggers_rebuild() { assert_eq!(results.len(), 3); // Deleted points should not appear for id in ids.iter().take(3) { - assert!(!results.contains(id)); + assert!(!results.iter().any(|(r, _)| r == id)); } } @@ -405,7 +405,7 @@ fn test_single_point_search() { .search(vec![0.0, 0.0], Similarity::Euclidean, 1) .unwrap(); assert_eq!(results.len(), 1); - assert_eq!(results[0], id); + assert_eq!(results[0].0, id); } #[test] @@ -425,7 +425,7 @@ fn test_duplicate_coordinates() { .unwrap(); assert_eq!(results.len(), 2); // Both id1 and id2 should be in results (both at distance 0) - assert!(results.contains(&id1) || results.contains(&id2)); + assert!(results.iter().any(|(r, _)| *r == id1 || *r == id2)); } #[test] @@ -441,7 +441,7 @@ fn test_negative_coordinates() { let results = tree .search(vec![-0.5, -0.5], Similarity::Euclidean, 1) .unwrap(); - assert_eq!(results[0], id1); + assert_eq!(results[0].0, id1); } #[test] @@ -493,8 +493,8 @@ fn create_test_vectors_3d() -> Vec { /// Helper to verify that two result sets are valid k-nearest neighbor results /// Both should return the k closest points (by distance), but may differ on tie-breaking fn verify_same_results( - tree_results: &[Uuid], - flat_results: &[Uuid], + tree_results: &[(Uuid, f32)], + flat_results: &[(Uuid, f32)], vectors: &[IndexedVector], query: &[f32], similarity: Similarity, @@ -510,6 +510,9 @@ fn verify_same_results( // Both should return at most k results assert!(tree_results.len() <= k); + let tree_results: Vec = tree_results.iter().map(|(id, _)| *id).collect(); + let flat_results: Vec = flat_results.iter().map(|(id, _)| *id).collect(); + // Get distances for all results let query_vec = query.to_vec(); let get_distance = |id: &Uuid| -> f32 { @@ -558,7 +561,7 @@ fn verify_same_results( get_distance(flat_results.last().unwrap()) }; - for id in tree_results { + for id in &tree_results { if !flat_set.contains(id) { // This ID is not in flat results, verify it's a tie let dist = get_distance(id); @@ -580,7 +583,7 @@ fn verify_same_results( get_distance(tree_results.last().unwrap()) }; - for id in flat_results { + for id in &flat_results { if !tree_set.contains(id) { let dist = get_distance(id); assert!( diff --git a/crates/index/src/lib.rs b/crates/index/src/lib.rs index 8b80cd9..48452e3 100644 --- a/crates/index/src/lib.rs +++ b/crates/index/src/lib.rs @@ -20,7 +20,7 @@ pub trait VectorIndex: Send + Sync + SerializableIndex { query_vector: DenseVector, similarity: Similarity, k: usize, - ) -> Result>; // Return a Vec of ids of closest vectors (length max k) + ) -> Result>; // Return a Vec of (id, distance) of closest vectors (length max k) fn search_with_ef( &self, @@ -28,7 +28,7 @@ pub trait VectorIndex: Send + Sync + SerializableIndex { similarity: Similarity, k: usize, _ef: Option, - ) -> Result> { + ) -> Result> { self.search(query_vector, similarity, k) } @@ -72,6 +72,9 @@ pub fn distance(a: &[f32], b: &[f32], dist_type: Similarity) -> f32 { norm_a += x * x; norm_b += y * y; } + if norm_a == 0.0 || norm_b == 0.0 { + return 1.0; + } 1.0 - dot / (norm_a.sqrt() * norm_b.sqrt()) } } diff --git a/crates/tui/src/app/events.rs b/crates/tui/src/app/events.rs index e7d90a9..c2bf708 100644 --- a/crates/tui/src/app/events.rs +++ b/crates/tui/src/app/events.rs @@ -326,7 +326,7 @@ fn execute_modal_action(app: &mut App) -> io::Result<()> { app.vector_list_post_restore = false; app.vector_list_selected_index = 0; - for id in ids.into_iter() { + for (id, _) in ids.into_iter() { if let Some(point) = db.get(id).map_err(to_io)? && let Some(vector) = point.vector {