added sample for getting around 2mb limit

Author: daniel-sanche

datastore/cloud-client/vector_search.py

@@ -116,4 +116,35 @@ def vector_search_distance_threshold(db):
     for entity in vector_query.fetch():
         print(f"{entity.id}")
     # [END datastore_vector_search_distance_threshold]
-    return vector_query
+    return vector_query
+
+
+def vector_search_large_query(db):
+    # [START datastore_vector_search_large_query]
+    from google.cloud.datastore.vector import DistanceMeasure
+    from google.cloud.datastore.vector import Vector
+    from google.cloud.datastore.vector import FindNearest
+
+    # first, perform a vector search query retrieving just the keys
+    vector_query = db.query(
+        kind="coffee-beans",
+        find_nearest=FindNearest(
+            vector_property="embedding_field",
+            query_vector=Vector([3.0, 1.0, 2.0]),
+            distance_measure=DistanceMeasure.EUCLIDEAN,
+            limit=100,
+            distance_result_property="vector_distance",
+        )
+    )
+    vector_query.keys_only()
+    vector_results = list(vector_query.fetch())
+    key_list = [entity.key for entity in vector_results]
+    # next, perfrom a second query for the remaining data
+    full_results = db.get_multi(key_list)
+    # combine and print results
+    vector_map = {entity.key: entity for entity in vector_results}
+    full_map = {entity.key: entity for entity in full_results}
+    for key in key_list:
+        print(f"distance: {vector_map[key]['vector_distance']} entity: {full_map[key]}")
+    # [END datastore_vector_search_large_query]
+    return key_list, vector_results, full_results

datastore/cloud-client/vector_search_test.py

@@ -106,4 +106,29 @@ def test_vector_search_distance_threshold(db):
 
     assert len(results) == 2
     assert results[0].key.name == "Liberica"
-    assert results[1].key.name == "Robusta"
+    assert results[1].key.name == "Robusta"
+
+def test_vector_search_large_query(db):
+    key_list, vector_results, full_results = vector_search_large_query(db)
+    assert len(key_list) == 4
+    # each list should have same number of elements
+    assert len(key_list) == len(vector_results)
+    assert len(key_list) == len(full_results)
+    # should all have the same keys
+    vector_map = {entity.key: entity for entity in vector_results}
+    full_map = {entity.key: entity for entity in full_results}
+    for key in key_list:
+        assert key in vector_map.keys()
+        assert key in full_map.keys()
+    # vector_results should just contain key and distance
+    for entity in vector_results:
+        assert entity.key is not None
+        assert entity["vector_distance"] is not None
+        with pytest.raises(KeyError):
+            entity["embedding_field"]
+    # full_results should have other fields, but no vector_distance
+    for entity in full_results:
+        assert entity.key is not None
+        assert isinstance(entity["embedding_field"], Vector)
+        with pytest.raises(KeyError):
+            entity["vector_distance"]