feat: add SearchNearestEntities rpc to FeatureOnlineStoreService in a…

…iplatform v1

PiperOrigin-RevId: 604763557

google/cloud/aiplatform/v1/aiplatform_v1.yaml

@@ -630,6 +630,10 @@ authentication:
     oauth:
       canonical_scopes: |-
         https://www.googleapis.com/auth/cloud-platform
+  - selector: google.cloud.aiplatform.v1.FeatureOnlineStoreService.SearchNearestEntities
+    oauth:
+      canonical_scopes: |-
+        https://www.googleapis.com/auth/cloud-platform
   - selector: 'google.cloud.aiplatform.v1.FeatureRegistryService.*'
     oauth:
       canonical_scopes: |-

google/cloud/aiplatform/v1/feature_online_store_service.proto

@@ -46,6 +46,17 @@ service FeatureOnlineStoreService {
     };
     option (google.api.method_signature) = "feature_view, data_key";
   }
+
+  // Search the nearest entities under a FeatureView.
+  // Search only works for indexable feature view; if a feature view isn't
+  // indexable, returns Invalid argument response.
+  rpc SearchNearestEntities(SearchNearestEntitiesRequest)
+      returns (SearchNearestEntitiesResponse) {
+    option (google.api.http) = {
+      post: "/v1/{feature_view=projects/*/locations/*/featureOnlineStores/*/featureViews/*}:searchNearestEntities"
+      body: "*"
+    };
+  }
 }
 
 // Format of the data in the Feature View.
@@ -120,3 +131,126 @@ message FetchFeatureValuesResponse {
     google.protobuf.Struct proto_struct = 2;
   }
 }
+
+// A query to find a number of similar entities.
+message NearestNeighborQuery {
+  // The embedding vector.
+  message Embedding {
+    // Optional. Individual value in the embedding.
+    repeated float value = 1 [(google.api.field_behavior) = OPTIONAL];
+  }
+
+  // String filter is used to search a subset of the entities by using boolean
+  // rules on string columns.
+  // For example: if a query specifies string filter
+  // with 'name = color, allow_tokens = {red, blue}, deny_tokens = {purple}','
+  // then that query will match entities that are red or blue, but if those
+  // points are also purple, then they will be excluded even if they are
+  // red/blue. Only string filter is supported for now, numeric filter will be
+  // supported in the near future.
+  message StringFilter {
+    // Required. Column names in BigQuery that used as filters.
+    string name = 1 [(google.api.field_behavior) = REQUIRED];
+
+    // Optional. The allowed tokens.
+    repeated string allow_tokens = 2 [(google.api.field_behavior) = OPTIONAL];
+
+    // Optional. The denied tokens.
+    repeated string deny_tokens = 3 [(google.api.field_behavior) = OPTIONAL];
+  }
+
+  // Parameters that can be overrided in each query to tune query latency and
+  // recall.
+  message Parameters {
+    // Optional. The number of neighbors to find via approximate search before
+    // exact reordering is performed; if set, this value must be >
+    // neighbor_count.
+    int32 approximate_neighbor_candidates = 1
+        [(google.api.field_behavior) = OPTIONAL];
+
+    // Optional. The fraction of the number of leaves to search, set at query
+    // time allows user to tune search performance. This value increase result
+    // in both search accuracy and latency increase. The value should be between
+    // 0.0 and 1.0.
+    double leaf_nodes_search_fraction = 2
+        [(google.api.field_behavior) = OPTIONAL];
+  }
+
+  oneof instance {
+    // Optional. The entity id whose similar entities should be searched for.
+    // If embedding is set, search will use embedding instead of
+    // entity_id.
+    string entity_id = 1 [(google.api.field_behavior) = OPTIONAL];
+
+    // Optional. The embedding vector that be used for similar search.
+    Embedding embedding = 2 [(google.api.field_behavior) = OPTIONAL];
+  }
+
+  // Optional. The number of similar entities to be retrieved from feature view
+  // for each query.
+  int32 neighbor_count = 3 [(google.api.field_behavior) = OPTIONAL];
+
+  // Optional. The list of string filters.
+  repeated StringFilter string_filters = 4
+      [(google.api.field_behavior) = OPTIONAL];
+
+  // Optional. Crowding is a constraint on a neighbor list produced by nearest
+  // neighbor search requiring that no more than
+  // sper_crowding_attribute_neighbor_count of the k neighbors returned have the
+  // same value of crowding_attribute. It's used for improving result diversity.
+  int32 per_crowding_attribute_neighbor_count = 5
+      [(google.api.field_behavior) = OPTIONAL];
+
+  // Optional. Parameters that can be set to tune query on the fly.
+  Parameters parameters = 7 [(google.api.field_behavior) = OPTIONAL];
+}
+
+// The request message for
+// [FeatureOnlineStoreService.SearchNearestEntities][google.cloud.aiplatform.v1.FeatureOnlineStoreService.SearchNearestEntities].
+message SearchNearestEntitiesRequest {
+  // Required. FeatureView resource format
+  // `projects/{project}/locations/{location}/featureOnlineStores/{featureOnlineStore}/featureViews/{featureView}`
+  string feature_view = 1 [
+    (google.api.field_behavior) = REQUIRED,
+    (google.api.resource_reference) = {
+      type: "aiplatform.googleapis.com/FeatureView"
+    }
+  ];
+
+  // Required. The query.
+  NearestNeighborQuery query = 2 [(google.api.field_behavior) = REQUIRED];
+
+  // Optional. If set to true, the full entities (including all vector values
+  // and metadata) of the nearest neighbors are returned; otherwise only entity
+  // id of the nearest neighbors will be returned. Note that returning full
+  // entities will significantly increase the latency and cost of the query.
+  bool return_full_entity = 3 [(google.api.field_behavior) = OPTIONAL];
+}
+
+// Nearest neighbors for one query.
+message NearestNeighbors {
+  // A neighbor of the query vector.
+  message Neighbor {
+    // The id of the similar entity.
+    string entity_id = 1;
+
+    // The distance between the neighbor and the query vector.
+    double distance = 2;
+
+    // The attributes of the neighbor, e.g. filters, crowding and metadata
+    // Note that full entities are returned only when "return_full_entity"
+    // is set to true. Otherwise, only the "entity_id" and "distance" fields
+    // are populated.
+    FetchFeatureValuesResponse entity_key_values = 3;
+  }
+
+  // All its neighbors.
+  repeated Neighbor neighbors = 1;
+}
+
+// Response message for
+// [FeatureOnlineStoreService.SearchNearestEntities][google.cloud.aiplatform.v1.FeatureOnlineStoreService.SearchNearestEntities]
+message SearchNearestEntitiesResponse {
+  // The nearest neighbors of the query entity.
+  NearestNeighbors nearest_neighbors = 1;
+}