Skip to content

Commit 60dd19d

Browse files
cindy-pengcindy-peng
committed
Add sample code for Java datastore vector search
1 parent 9d79efc commit 60dd19d

File tree

8 files changed

+412
-1
lines changed

8 files changed

+412
-1
lines changed

‎samples/snippets/src/main/java/com/example/datastore/filters/OrderFieldsQuery.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717
package com.example.datastore.filters;
1818

1919
// sample-metadata:
20-
// title: Queries with order fileds
20+
// title: Queries with order fields
2121
// description: The following query order properties
2222
// in the decreasing order of query constraint selectivity.
2323

‎samples/snippets/src/main/java/com/example/datastore/vectorsearch/StoreVectors.java‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
/*
2+
* Copyright 2024 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.example.datastore.vectorsearch;
18+
19+
// [START datastore_store_vectors]
20+
21+
import com.google.cloud.datastore.Datastore;
22+
import com.google.cloud.datastore.DatastoreOptions;
23+
import com.google.cloud.datastore.Key;
24+
import com.google.cloud.datastore.Entity;
25+
import com.google.cloud.datastore.VectorValue;
26+
27+
public class StoreVectors {
28+
public static void invoke() throws Exception {
29+
// Instantiates a client
30+
Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
31+
32+
// The Cloud Datastore key for the new entity
33+
Key key = datastore.newKeyFactory().setKind("CoffeeBean").newKey("Kahawa");
34+
35+
// Prepares the entity with a vector embedding
36+
Entity entity =
37+
Entity.newBuilder(key)
38+
.set("name", "Kahawa")
39+
.set("description", "Information about the Kahawa coffee beans.")
40+
.set("roast", "dark")
41+
.set("embedding_field", VectorValue.newBuilder(1.0, 7.0, 11.1).build())
42+
.build();
43+
44+
// Saves the entity
45+
datastore.put(entity);
46+
System.out.printf("Saved %s: %s%n", entity.getKey().getName(), entity.getString("description"));
47+
48+
// Retrieve entity
49+
Entity retrieved = datastore.get(key);
50+
System.out.printf("Retrieved %s: %s%n", key.getName(), retrieved.getString("description"));
51+
}
52+
}
53+
// [END datastore_store_vectors]

‎samples/snippets/src/main/java/com/example/datastore/vectorsearch/VectorSearchBasic.java‎

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
1+
/*
2+
* Copyright 2024 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.example.datastore.vectorsearch;
18+
19+
// [START datastore_vector_search_basic]
20+
21+
import com.google.cloud.datastore.Datastore;
22+
import com.google.cloud.datastore.DatastoreOptions;
23+
import com.google.cloud.datastore.Entity;
24+
import com.google.cloud.datastore.Query;
25+
import com.google.cloud.datastore.QueryResults;
26+
import com.google.cloud.datastore.VectorValue;
27+
import com.google.cloud.datastore.FindNearest;
28+
29+
public class VectorSearchBasic {
30+
public static void invoke() throws Exception {
31+
// Instantiates a client
32+
Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
33+
34+
// Create vector search query
35+
Query<Entity> vectorSearchQuery =
36+
Query.newEntityQueryBuilder()
37+
.setKind("CoffeeBean")
38+
.setFindNearest(new FindNearest(
39+
"embedding_field",
40+
VectorValue.newBuilder(1, 9, 11.1).build(),
41+
FindNearest.DistanceMeasure.DOT_PRODUCT,
42+
3))
43+
.build();
44+
45+
// Execute vector search query
46+
QueryResults<Entity> results = datastore.run(vectorSearchQuery);
47+
48+
if (!results.hasNext()) {
49+
throw new Exception("query yielded no results");
50+
}
51+
52+
while (results.hasNext()) {
53+
Entity entity = results.next();
54+
System.out.printf("Entity: %s%n", entity.getKey().getName());
55+
}
56+
}
57+
}
58+
// [END datastore_vector_search_basic]

‎samples/snippets/src/main/java/com/example/datastore/vectorsearch/VectorSearchDistanceResultProperty.java‎

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
1+
/*
2+
* Copyright 2024 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.example.datastore.vectorsearch;
18+
19+
// [START datastore_vector_search_distance_result_property]
20+
21+
import com.google.cloud.datastore.Datastore;
22+
import com.google.cloud.datastore.DatastoreOptions;
23+
import com.google.cloud.datastore.Entity;
24+
import com.google.cloud.datastore.Query;
25+
import com.google.cloud.datastore.QueryResults;
26+
import com.google.cloud.datastore.VectorValue;
27+
import com.google.cloud.datastore.FindNearest;
28+
29+
public class VectorSearchDistanceResultProperty {
30+
public static void invoke() throws Exception {
31+
// Instantiates a client
32+
Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
33+
34+
// Create vector search query with distance result property
35+
Query<Entity> vectorSearchQuery =
36+
Query.newEntityQueryBuilder()
37+
.setKind("CoffeeBean")
38+
.setFindNearest(new FindNearest(
39+
"embedding_field",
40+
VectorValue.newBuilder(1, 9, 11.1).build(),
41+
FindNearest.DistanceMeasure.DOT_PRODUCT,
42+
3, "vector_distance"))
43+
.build();
44+
45+
// Execute vector search query
46+
QueryResults<Entity> results = datastore.run(vectorSearchQuery);
47+
48+
if (!results.hasNext()) {
49+
throw new Exception("query yielded no results");
50+
}
51+
52+
while (results.hasNext()) {
53+
Entity entity = results.next();
54+
System.out.printf("Entity: %s, Distance: %s%n", entity.getKey().getName(), entity.getDouble("vector_distance"));
55+
}
56+
}
57+
}
58+
// [END datastore_vector_search_distance_result_property]

‎samples/snippets/src/main/java/com/example/datastore/vectorsearch/VectorSearchDistanceResultPropertyProjection.java‎

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
/*
2+
* Copyright 2024 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.example.datastore.vectorsearch;
18+
19+
// [START datastore_vector_search_distance_result_property_projection]
20+
21+
import com.google.cloud.datastore.Datastore;
22+
import com.google.cloud.datastore.DatastoreOptions;
23+
import com.google.cloud.datastore.Query;
24+
import com.google.cloud.datastore.QueryResults;
25+
import com.google.cloud.datastore.ProjectionEntity;
26+
import com.google.cloud.datastore.VectorValue;
27+
import com.google.cloud.datastore.FindNearest;
28+
29+
public class VectorSearchDistanceResultPropertyProjection {
30+
public static void invoke() throws Exception {
31+
// Instantiates a client
32+
Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
33+
34+
// Create vector search query with projection
35+
Query<ProjectionEntity> vectorSearchQuery =
36+
Query.newProjectionEntityQueryBuilder()
37+
.setKind("CoffeeBean")
38+
.setFindNearest(new FindNearest(
39+
"embedding_field",
40+
VectorValue.newBuilder(1, 9, 11.1).build(),
41+
FindNearest.DistanceMeasure.DOT_PRODUCT,
42+
3, "vector_distance"))
43+
.setProjection("vector_distance")
44+
.build();
45+
46+
// Execute vector search query
47+
QueryResults<ProjectionEntity> results = datastore.run(vectorSearchQuery);
48+
49+
if (!results.hasNext()) {
50+
throw new Exception("query yielded no results");
51+
}
52+
53+
while (results.hasNext()) {
54+
ProjectionEntity entity = results.next();
55+
System.out.printf("Entity: %s, Distance: %s%n", entity.getKey().getName(), entity.getDouble("vector_distance"));
56+
}
57+
}
58+
}
59+
// [END datastore_vector_search_distance_result_property_projection]

‎samples/snippets/src/main/java/com/example/datastore/vectorsearch/VectorSearchDistanceThreshold.java‎

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
1+
/*
2+
* Copyright 2024 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.example.datastore.vectorsearch;
18+
19+
// [START datastore_vector_search_distance_threshold]
20+
21+
import com.google.cloud.datastore.Datastore;
22+
import com.google.cloud.datastore.DatastoreOptions;
23+
import com.google.cloud.datastore.Entity;
24+
import com.google.cloud.datastore.Query;
25+
import com.google.cloud.datastore.QueryResults;
26+
import com.google.cloud.datastore.VectorValue;
27+
import com.google.cloud.datastore.FindNearest;
28+
29+
public class VectorSearchDistanceThreshold {
30+
public static void invoke() throws Exception {
31+
// Instantiates a client
32+
Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
33+
34+
// Create vector search query with distance threshold
35+
Query<Entity> vectorSearchQuery =
36+
Query.newEntityQueryBuilder()
37+
.setKind("CoffeeBean")
38+
.setFindNearest(new FindNearest(
39+
"embedding_field",
40+
VectorValue.newBuilder(1, 9, 11.1).build(),
41+
FindNearest.DistanceMeasure.EUCLIDEAN,
42+
3, "vector_distance", 2.0))
43+
.build();
44+
45+
// Execute vector search query
46+
QueryResults<Entity> results = datastore.run(vectorSearchQuery);
47+
48+
if (!results.hasNext()) {
49+
throw new Exception("query yielded no results");
50+
}
51+
52+
while (results.hasNext()) {
53+
Entity entity = results.next();
54+
System.out.printf("Entity: %s, Distance: %s%n", entity.getKey().getName(), entity.getDouble("vector_distance"));
55+
}
56+
}
57+
}
58+
// [END datastore_vector_search_distance_threshold]

‎samples/snippets/src/main/java/com/example/datastore/vectorsearch/VectorSearchLargeResponse.java‎

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,65 @@
1+
/*
2+
* Copyright 2024 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.example.datastore.vectorsearch;
18+
19+
// [START datastore_vector_search_large_response]
20+
21+
import com.google.common.collect.Iterators;
22+
import java.util.Iterator;
23+
import com.google.cloud.datastore.Datastore;
24+
import com.google.cloud.datastore.DatastoreOptions;
25+
import com.google.cloud.datastore.Entity;
26+
import com.google.cloud.datastore.Query;
27+
import com.google.cloud.datastore.QueryResults;
28+
import com.google.cloud.datastore.Key;
29+
import com.google.cloud.datastore.VectorValue;
30+
import com.google.cloud.datastore.FindNearest;
31+
32+
public class VectorSearchLargeResponse {
33+
public static void invoke() throws Exception {
34+
// Instantiates a client
35+
Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
36+
37+
// Create a keys-only vector search query
38+
Query<Key> vectorSearchKeyOnlyQuery =
39+
Query.newKeyQueryBuilder()
40+
.setKind("CoffeeBean")
41+
.setFindNearest(new FindNearest(
42+
"embedding_field",
43+
VectorValue.newBuilder(1, 9, 11.1).build(),
44+
FindNearest.DistanceMeasure.EUCLIDEAN,
45+
3, "vector_distance", 2.0))
46+
.build();
47+
48+
QueryResults<Key> keyResults = datastore.run(vectorSearchKeyOnlyQuery);
49+
Key[] keys = Iterators.toArray(keyResults, Key.class);
50+
51+
// Next, perform a second query for the remaining data
52+
Iterator<Entity> entities = datastore.get(keys);
53+
54+
if (!keyResults.hasNext()) {
55+
throw new Exception("query yielded no results");
56+
}
57+
58+
// Combine and print results
59+
while (keyResults.hasNext()) {
60+
Key keyResult = keyResults.next();
61+
System.out.printf("Entity: %s, Distance: %s%n", keyResult.getName(), entities.next().getDouble("vector_distance"));
62+
}
63+
}
64+
}
65+
// [END datastore_vector_search_large_response]

0 commit comments

Comments
 (0)