Skip to content
This repository was archived by the owner on Mar 2, 2026. It is now read-only.

Commit 5c2192d

Browse files
NickChittleNickChittle
authored
feat: support returning computed distance and set distance thresholds on VectorQueries (#960)
1 parent 53b8aab commit 5c2192d

10 files changed

Lines changed: 768 additions & 37 deletions

File tree

‎google/cloud/firestore_v1/async_query.py‎

Lines changed: 13 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -230,17 +230,25 @@ def find_nearest(
230230
query_vector: Vector,
231231
limit: int,
232232
distance_measure: DistanceMeasure,
233+
*,
234+
distance_result_field: Optional[str] = None,
235+
distance_threshold: Optional[float] = None,
233236
) -> AsyncVectorQuery:
234237
"""
235238
Finds the closest vector embeddings to the given query vector.
236239
237240
Args:
238-
vector_field(str): An indexed vector field to search upon. Only documents which contain
241+
vector_field (str): An indexed vector field to search upon. Only documents which contain
239242
vectors whose dimensionality match the query_vector can be returned.
240-
query_vector(Vector): The query vector that we are searching on. Must be a vector of no more
243+
query_vector (Vector): The query vector that we are searching on. Must be a vector of no more
241244
than 2048 dimensions.
242245
limit (int): The number of nearest neighbors to return. Must be a positive integer of no more than 1000.
243-
distance_measure(:class:`DistanceMeasure`): The Distance Measure to use.
246+
distance_measure (:class:`DistanceMeasure`): The Distance Measure to use.
247+
distance_result_field (Optional[str]):
248+
Name of the field to output the result of the vector distance
249+
calculation. If unset then the distance will not be returned.
250+
distance_threshold (Optional[float]):
251+
A threshold for which no less similar documents will be returned.
244252
245253
Returns:
246254
:class`~firestore_v1.vector_query.VectorQuery`: the vector query.
@@ -250,6 +258,8 @@ def find_nearest(
250258
query_vector=query_vector,
251259
limit=limit,
252260
distance_measure=distance_measure,
261+
distance_result_field=distance_result_field,
262+
distance_threshold=distance_threshold,
253263
)
254264

255265
def count(

‎google/cloud/firestore_v1/base_collection.py‎

Lines changed: 16 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -550,23 +550,35 @@ def find_nearest(
550550
query_vector: Vector,
551551
limit: int,
552552
distance_measure: DistanceMeasure,
553+
*,
554+
distance_result_field: Optional[str] = None,
555+
distance_threshold: Optional[float] = None,
553556
) -> VectorQuery:
554557
"""
555558
Finds the closest vector embeddings to the given query vector.
556559
557560
Args:
558-
vector_field(str): An indexed vector field to search upon. Only documents which contain
561+
vector_field (str): An indexed vector field to search upon. Only documents which contain
559562
vectors whose dimensionality match the query_vector can be returned.
560-
query_vector(Vector): The query vector that we are searching on. Must be a vector of no more
563+
query_vector (Vector): The query vector that we are searching on. Must be a vector of no more
561564
than 2048 dimensions.
562565
limit (int): The number of nearest neighbors to return. Must be a positive integer of no more than 1000.
563-
distance_measure(:class:`DistanceMeasure`): The Distance Measure to use.
566+
distance_measure (:class:`DistanceMeasure`): The Distance Measure to use.
567+
distance_result_field (Optional[str]):
568+
Name of the field to output the result of the vector distance calculation
569+
distance_threshold (Optional[float]):
570+
A threshold for which no less similar documents will be returned.
564571
565572
Returns:
566573
:class`~firestore_v1.vector_query.VectorQuery`: the vector query.
567574
"""
568575
return self._vector_query().find_nearest(
569-
vector_field, query_vector, limit, distance_measure
576+
vector_field,
577+
query_vector,
578+
limit,
579+
distance_measure,
580+
distance_result_field=distance_result_field,
581+
distance_threshold=distance_threshold,
570582
)
571583

572584

‎google/cloud/firestore_v1/base_query.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -982,6 +982,9 @@ def find_nearest(
982982
query_vector: Vector,
983983
limit: int,
984984
distance_measure: DistanceMeasure,
985+
*,
986+
distance_result_field: Optional[str] = None,
987+
distance_threshold: Optional[float] = None,
985988
) -> BaseVectorQuery:
986989
raise NotImplementedError
987990

‎google/cloud/firestore_v1/base_vector_query.py‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,8 @@ def __init__(self, nested_query) -> None:
4545
self._query_vector: Optional[Vector] = None
4646
self._limit: Optional[int] = None
4747
self._distance_measure: Optional[DistanceMeasure] = None
48+
self._distance_result_field: Optional[str] = None
49+
self._distance_threshold: Optional[float] = None
4850

4951
@property
5052
def _client(self):
@@ -69,6 +71,11 @@ def _to_protobuf(self) -> query.StructuredQuery:
6971
else:
7072
raise ValueError("Invalid distance_measure")
7173

74+
# Coerce ints to floats as required by the protobuf.
75+
distance_threshold_proto = None
76+
if self._distance_threshold is not None:
77+
distance_threshold_proto = float(self._distance_threshold)
78+
7279
pb = self._nested_query._to_protobuf()
7380
pb.find_nearest = query.StructuredQuery.FindNearest(
7481
vector_field=query.StructuredQuery.FieldReference(
@@ -77,6 +84,8 @@ def _to_protobuf(self) -> query.StructuredQuery:
7784
query_vector=_helpers.encode_value(self._query_vector),
7885
distance_measure=distance_measure_proto,
7986
limit=self._limit,
87+
distance_result_field=self._distance_result_field,
88+
distance_threshold=distance_threshold_proto,
8089
)
8190
return pb
8291

@@ -111,12 +120,17 @@ def find_nearest(
111120
query_vector: Vector,
112121
limit: int,
113122
distance_measure: DistanceMeasure,
123+
*,
124+
distance_result_field: Optional[str] = None,
125+
distance_threshold: Optional[float] = None,
114126
):
115127
"""Finds the closest vector embeddings to the given query vector."""
116128
self._vector_field = vector_field
117129
self._query_vector = query_vector
118130
self._limit = limit
119131
self._distance_measure = distance_measure
132+
self._distance_result_field = distance_result_field
133+
self._distance_threshold = distance_threshold
120134
return self
121135

122136
def stream(

‎google/cloud/firestore_v1/query.py‎

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -251,17 +251,26 @@ def find_nearest(
251251
query_vector: Vector,
252252
limit: int,
253253
distance_measure: DistanceMeasure,
254+
*,
255+
distance_result_field: Optional[str] = None,
256+
distance_threshold: Optional[float] = None,
254257
) -> Type["firestore_v1.vector_query.VectorQuery"]:
255258
"""
256259
Finds the closest vector embeddings to the given query vector.
257260
258261
Args:
259-
vector_field(str): An indexed vector field to search upon. Only documents which contain
262+
vector_field (str): An indexed vector field to search upon. Only documents which contain
260263
vectors whose dimensionality match the query_vector can be returned.
261-
query_vector(Vector): The query vector that we are searching on. Must be a vector of no more
264+
query_vector (Vector): The query vector that we are searching on. Must be a vector of no more
262265
than 2048 dimensions.
263266
limit (int): The number of nearest neighbors to return. Must be a positive integer of no more than 1000.
264-
distance_measure(:class:`DistanceMeasure`): The Distance Measure to use.
267+
distance_measure (:class:`DistanceMeasure`): The Distance Measure to use.
268+
distance_result_field (Optional[str]):
269+
Name of the field to output the result of the vector distance
270+
calculation. If unset then the distance will not be returned.
271+
distance_threshold (Optional[float]):
272+
A threshold for which no less similar documents will be returned.
273+
265274
266275
Returns:
267276
:class`~firestore_v1.vector_query.VectorQuery`: the vector query.
@@ -271,6 +280,8 @@ def find_nearest(
271280
query_vector=query_vector,
272281
limit=limit,
273282
distance_measure=distance_measure,
283+
distance_result_field=distance_result_field,
284+
distance_threshold=distance_threshold,
274285
)
275286

276287
def count(

0 commit comments

Comments
 (0)