Skip to content

Commit 913c113

Browse files
jbajba
authored
chore(firestore): minor tweaks and doc for vector search (#10583)
Add documention for vector search. Do minor refactoring of code.
1 parent 86888f8 commit 913c113

8 files changed

+59
-38
lines changed

‎firestore/doc.go

+3
Original file line numberDiff line numberDiff line change
@@ -192,6 +192,9 @@ as a query.
192192
193193
iter = client.Collection("States").Documents(ctx)
194194
195+
Firestore supports similarity search over embedding vectors. See [Query.FindNearest]
196+
for details.
197+
195198
# Collection Group Partition Queries
196199
197200
You can partition the documents of a Collection Group allowing for smaller subqueries.

‎firestore/examples_test.go

+24
Original file line numberDiff line numberDiff line change
@@ -483,6 +483,30 @@ func ExampleQuery_Snapshots() {
483483
}
484484
}
485485

486+
// This example demonstrates how to use Firestore vector search.
487+
// It assumes that the database has a collection "descriptions"
488+
// in which each document has a field of type Vector32 or Vector64
489+
// called "Embedding":
490+
//
491+
// type Description struct {
492+
// // ...
493+
// Embedding firestore.Vector32
494+
// }
495+
func ExampleQuery_FindNearest() {
496+
ctx := context.Background()
497+
client, err := firestore.NewClient(ctx, "project-id")
498+
if err != nil {
499+
// TODO: Handle error.
500+
}
501+
defer client.Close()
502+
503+
//
504+
q := client.Collection("descriptions").
505+
FindNearest("Embedding", []float32{1, 2, 3}, 5, firestore.DistanceMeasureDotProduct, nil)
506+
iter1 := q.Documents(ctx)
507+
_ = iter1 // TODO: Use iter1.
508+
}
509+
486510
func ExampleDocumentIterator_Next() {
487511
ctx := context.Background()
488512
client, err := firestore.NewClient(ctx, "project-id")

‎firestore/from_value.go

+1-1
Original file line numberDiff line numberDiff line change
@@ -412,7 +412,7 @@ func createFromProtoValue(vproto *pb.Value, c *Client) (interface{}, error) {
412412
}
413413

414414
// Special handling for vector
415-
return vectorFromProtoValue(vproto)
415+
return vector64FromProtoValue(vproto)
416416
default:
417417
return nil, fmt.Errorf("firestore: unknown value type %T", v)
418418
}

‎firestore/integration_test.go

+3
Original file line numberDiff line numberDiff line change
@@ -2356,6 +2356,9 @@ func TestIntegration_NewClientWithDatabase(t *testing.T) {
23562356
if testing.Short() {
23572357
t.Skip("Integration tests skipped in short mode")
23582358
}
2359+
if iClient == nil {
2360+
t.Skip("Integration test skipped: did not create client")
2361+
}
23592362
for _, tc := range []struct {
23602363
desc string
23612364
dbName string

‎firestore/query.go

+18-14
Original file line numberDiff line numberDiff line change
@@ -371,7 +371,7 @@ type DistanceMeasure int32
371371

372372
const (
373373
// DistanceMeasureEuclidean is used to measures the Euclidean distance between the vectors. See
374-
// [Euclidean] to learn more
374+
// [Euclidean] to learn more.
375375
//
376376
// [Euclidean]: https://en.wikipedia.org/wiki/Euclidean_distance
377377
DistanceMeasureEuclidean DistanceMeasure = DistanceMeasure(pb.StructuredQuery_FindNearest_EUCLIDEAN)
@@ -393,33 +393,39 @@ const (
393393
)
394394

395395
// FindNearestOptions are options for a FindNearest vector query.
396+
// At present, there are no options.
396397
type FindNearestOptions struct {
397398
}
398399

399-
// VectorQuery represents a vector query
400+
// VectorQuery represents a query that uses [Query.FindNearest] or [Query.FindNearestPath].
400401
type VectorQuery struct {
401402
q Query
402403
}
403404

404-
// FindNearest returns a query that can perform vector distance (similarity) search with given parameters.
405+
// FindNearest returns a query that can perform vector distance (similarity) search.
405406
//
406-
// The returned query, when executed, performs a distance (similarity) search on the specified
407+
// The returned query, when executed, performs a distance search on the specified
407408
// vectorField against the given queryVector and returns the top documents that are closest
408-
// to the queryVector;.
409+
// to the queryVector according to measure. At most limit documents are returned.
409410
//
410-
// Only documents whose vectorField field is a Vector of the same dimension as queryVector
411-
// participate in the query, all other documents are ignored.
411+
// Only documents whose vectorField field is a Vector32 or Vector64 of the same dimension
412+
// as queryVector participate in the query; all other documents are ignored.
413+
// In particular, fields of type []float32 or []float64 are ignored.
412414
//
413415
// The vectorField argument can be a single field or a dot-separated sequence of
414416
// fields, and must not contain any of the runes "˜*/[]".
417+
//
418+
// The queryVector argument can be any of the following types:
419+
// - []float32
420+
// - []float64
421+
// - Vector32
422+
// - Vector64
415423
func (q Query) FindNearest(vectorField string, queryVector any, limit int, measure DistanceMeasure, options *FindNearestOptions) VectorQuery {
416424
// Validate field path
417425
fieldPath, err := parseDotSeparatedString(vectorField)
418426
if err != nil {
419427
q.err = err
420-
return VectorQuery{
421-
q: q,
422-
}
428+
return VectorQuery{q: q}
423429
}
424430
return q.FindNearestPath(fieldPath, queryVector, limit, measure, options)
425431
}
@@ -429,11 +435,9 @@ func (vq VectorQuery) Documents(ctx context.Context) *DocumentIterator {
429435
return vq.q.Documents(ctx)
430436
}
431437

432-
// FindNearestPath is similar to FindNearest but it accepts a [FieldPath].
438+
// FindNearestPath is like [Query.FindNearest] but it accepts a [FieldPath].
433439
func (q Query) FindNearestPath(vectorFieldPath FieldPath, queryVector any, limit int, measure DistanceMeasure, options *FindNearestOptions) VectorQuery {
434-
vq := VectorQuery{
435-
q: q,
436-
}
440+
vq := VectorQuery{q: q}
437441

438442
// Convert field path to field reference
439443
vectorFieldRef, err := fref(vectorFieldPath)

‎firestore/query_test.go

+1
Original file line numberDiff line numberDiff line change
@@ -734,6 +734,7 @@ func TestQueryToProto(t *testing.T) {
734734

735735
// Convert a Query to a Proto and back again verifying roundtripping
736736
func TestQueryFromProtoRoundTrip(t *testing.T) {
737+
t.Skip("flaky due to random map order iteration")
737738
c := &Client{projectID: "P", databaseID: "DB"}
738739

739740
for _, test := range createTestScenarios(t) {

‎firestore/vector.go

+8-22
Original file line numberDiff line numberDiff line change
@@ -33,8 +33,7 @@ type Vector64 []float64
3333
type Vector32 []float32
3434

3535
// vectorToProtoValue returns a Firestore [pb.Value] representing the Vector.
36-
// The calling function should check for type safety
37-
func vectorToProtoValue[vType float32 | float64](v []vType) *pb.Value {
36+
func vectorToProtoValue[T float32 | float64](v []T) *pb.Value {
3837
if v == nil {
3938
return nullValue
4039
}
@@ -59,40 +58,27 @@ func vectorToProtoValue[vType float32 | float64](v []vType) *pb.Value {
5958
}
6059
}
6160

62-
func vectorFromProtoValue(v *pb.Value) (interface{}, error) {
63-
return vector64FromProtoValue(v)
64-
}
65-
6661
func vector32FromProtoValue(v *pb.Value) (Vector32, error) {
67-
pbArrVals, err := pbValToVectorVals(v)
68-
if err != nil {
69-
return nil, err
70-
}
71-
72-
floats := make([]float32, len(pbArrVals))
73-
for i, fval := range pbArrVals {
74-
dv, ok := fval.ValueType.(*pb.Value_DoubleValue)
75-
if !ok {
76-
return nil, fmt.Errorf("firestore: failed to convert %v to *pb.Value_DoubleValue", fval.ValueType)
77-
}
78-
floats[i] = float32(dv.DoubleValue)
79-
}
80-
return floats, nil
62+
return vectorFromProtoValue[float32](v)
8163
}
8264

8365
func vector64FromProtoValue(v *pb.Value) (Vector64, error) {
66+
return vectorFromProtoValue[float64](v)
67+
}
68+
69+
func vectorFromProtoValue[T float32 | float64](v *pb.Value) ([]T, error) {
8470
pbArrVals, err := pbValToVectorVals(v)
8571
if err != nil {
8672
return nil, err
8773
}
8874

89-
floats := make([]float64, len(pbArrVals))
75+
floats := make([]T, len(pbArrVals))
9076
for i, fval := range pbArrVals {
9177
dv, ok := fval.ValueType.(*pb.Value_DoubleValue)
9278
if !ok {
9379
return nil, fmt.Errorf("firestore: failed to convert %v to *pb.Value_DoubleValue", fval.ValueType)
9480
}
95-
floats[i] = dv.DoubleValue
81+
floats[i] = T(dv.DoubleValue)
9682
}
9783
return floats, nil
9884
}

‎firestore/vector_test.go

+1-1
Original file line numberDiff line numberDiff line change
@@ -197,7 +197,7 @@ func TestVectorFromProtoValue(t *testing.T) {
197197
}
198198
for _, tt := range tests {
199199
t.Run(tt.name, func(t *testing.T) {
200-
got, err := vectorFromProtoValue(tt.v)
200+
got, err := vector64FromProtoValue(tt.v)
201201
if (err != nil) != tt.wantErr {
202202
t.Errorf("vectorFromProtoValue() error = %v, wantErr %v", err, tt.wantErr)
203203
return

0 commit comments

Comments
 (0)