Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(firestore): Adding vector search #10548

Merged
merged 12 commits into from
Jul 22, 2024

Conversation

bhshkh
Copy link
Contributor

@bhshkh bhshkh commented Jul 15, 2024

b/346827360

Python and Nodejs already have support for vector search. More details can be found here

This PR introduces the same functionality in Go client library.

Sample Usage:
Save documents with vector field:

type CoffeeBeans struct {
	EmbeddedField firestore.Vector64
}
.......
........
docRef1 := client.Collection(vectorCollection).Doc("Arabica")
docRef1.Create(ctx, CoffeeBeans{
EmbeddedField: firestore.Vector64([]float64{1.0, 2.0, 3.0}),
})

docRef2 := client.Collection(vectorCollection).Doc("Excelsa")
docRef2.Create(ctx, CoffeeBeans{
EmbeddedField: firestore.Vector64([]float64{4, 5, 6}),
})

docRef3 := client.Collection(vectorCollection).Doc("Liberica")
docRef3.Create(ctx, CoffeeBeans{
EmbeddedField: firestore.Vector64([]float64{400, 500, 600}),
})

FindNearest query:

func findNearest(client *firestore.Client, ctx context.Context) error {
	vectorQuery := client.Collection(vectorCollection).FindNearest("EmbeddedField", firestore.Vector64{1, 2, 3}, firestore.FindNearestOpts{
		Limit:   2,
		Measure: firestore.DistanceMeasureEuclidean,
	})

	iter := vectorQuery.Documents(ctx)
	gotDocs, _ := iter.GetAll()

	for _, doc := range gotDocs {
		bean := CoffeeBeans{}
		err := doc.DataTo(&bean)
		fmt.Printf("findNearest bean: %+v, err: %+v\n", bean, err)

                docMap := doc.Data()
	}
	return nil
}

The above query will return Excelsa and Arabica beans

Major changes:

  1. Added Vector64 type which is just an alias for []float64 and similarly Vector32
  2. Added VectorQuery type which is a wrapper for query but allows only the Documents method.
  3. Renamed variables in setReflectFromProtoValue function in firestore/from_value.go
  4. Added vector type handing in setReflectFromProtoValue method. This is used in the DataTo method demonstrated above where the user passes in a struct and the Firestore document is populated into the struct
  5. Added vector type handing in createFromProtoValue method. This is used in the Data method demonstrated above which returns the Firestore document in the form of a map.

@product-auto-label product-auto-label bot added the api: firestore Issues related to the Firestore API. label Jul 15, 2024
@bhshkh bhshkh requested a review from jba July 15, 2024 09:11
firestore/docref.go Outdated Show resolved Hide resolved
firestore/docref.go Outdated Show resolved Hide resolved
firestore/docref.go Outdated Show resolved Hide resolved
firestore/docref.go Outdated Show resolved Hide resolved
firestore/docref.go Outdated Show resolved Hide resolved
firestore/query.go Outdated Show resolved Hide resolved
firestore/query.go Outdated Show resolved Hide resolved
firestore/query.go Outdated Show resolved Hide resolved
firestore/to_value.go Outdated Show resolved Hide resolved
firestore/to_value.go Outdated Show resolved Hide resolved
@bhshkh bhshkh marked this pull request as ready for review July 16, 2024 18:25
@bhshkh bhshkh requested review from a team as code owners July 16, 2024 18:25
@jba
Copy link
Contributor

jba commented Jul 16, 2024

Thanks for the extensive tests!

firestore/from_value.go Outdated Show resolved Hide resolved
firestore/query.go Outdated Show resolved Hide resolved
firestore/from_value.go Outdated Show resolved Hide resolved
firestore/vector.go Outdated Show resolved Hide resolved
firestore/query.go Outdated Show resolved Hide resolved
firestore/query.go Outdated Show resolved Hide resolved
firestore/from_value.go Outdated Show resolved Hide resolved
firestore/vector.go Outdated Show resolved Hide resolved
Copy link
Contributor

@jba jba left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Very close! Some minor tweaks.

firestore/query.go Outdated Show resolved Hide resolved
firestore/query.go Outdated Show resolved Hide resolved
firestore/vector.go Outdated Show resolved Hide resolved
@bhshkh bhshkh enabled auto-merge (squash) July 22, 2024 17:18
@bhshkh bhshkh requested review from cindy-peng and jba July 22, 2024 17:19
firestore/query.go Outdated Show resolved Hide resolved
firestore/query.go Outdated Show resolved Hide resolved
@bhshkh bhshkh merged commit 5c0d6df into googleapis:main Jul 22, 2024
8 checks passed
@bhshkh bhshkh deleted the feature/fs-vector-search branch July 22, 2024 21:33
eliben pushed a commit to eliben/google-cloud-go that referenced this pull request Jul 23, 2024
* feat(firestore): Adding vector search

* feat(firestore): refactoring code

* feat(firestore): Resolving vet failures

* feat(firestore): Adding unit and integration tests

* feat(firestore): Fixing tests and refactoring code

* feat(firestore): Resolving vet failures

* feat(firestore): Refactoring code

* feat(firestore): Resolving review comments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
api: firestore Issues related to the Firestore API.
Projects
None yet
Development

Successfully merging this pull request may close these issues.

firestore: Add support for vector search in golang cloud firestore client
4 participants