feat: Support NOT_EQUAL for MatchingEngineIndexEndpoint `numeric_…

…restricts`. PiperOrigin-RevId: 623343218
googleapis · Apr 10, 2024 · aa918e3 · aa918e3
1 parent fa35b91
commit aa918e3
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 0 deletions.
diff --git a/google/cloud/aiplatform/matching_engine/matching_engine_index_endpoint.py b/google/cloud/aiplatform/matching_engine/matching_engine_index_endpoint.py
@@ -92,6 +92,7 @@ class NumericNamespace:
                 "EQUAL" for datapoints' values = query's value;
                 "GREATER_EQUAL" for datapoints' values >= query's value;
                 "GREATER" for datapoints' values > query's value;
+                "NOT_EQUAL" for datapoints' values != query's value
     """
 
     name: str

diff --git a/tests/unit/aiplatform/test_matching_engine_index_endpoint.py b/tests/unit/aiplatform/test_matching_engine_index_endpoint.py
@@ -240,11 +240,13 @@
     NumericNamespace(name="cost", value_double=0.3, op="EQUAL"),
     NumericNamespace(name="size", value_int=0, op="GREATER"),
     NumericNamespace(name="seconds", value_float=-20.5, op="LESS_EQUAL"),
+    NumericNamespace(name="duration", value_int=10, op="NOT_EQUAL"),
 ]
 _TEST_NUMERIC_NAMESPACE = [
     match_service_pb2.NumericNamespace(name="cost", value_double=0.3, op=3),
     match_service_pb2.NumericNamespace(name="size", value_int=0, op=5),
     match_service_pb2.NumericNamespace(name="seconds", value_float=-20.5, op=2),
+    match_service_pb2.NumericNamespace(name="duration", value_int=10, op="NOT_EQUAL"),
 ]
 _TEST_IDS = ["123", "456", "789"]
 _TEST_PER_CROWDING_ATTRIBUTE_NUM_NEIGHBOURS = 3
@@ -1421,6 +1423,9 @@ def test_index_public_endpoint_match_queries_with_numeric_filtering(
                             gca_index_v1beta1.IndexDatapoint.NumericRestriction(
                                 namespace="seconds", value_float=-20.5, op="LESS_EQUAL"
                             ),
+                            gca_index_v1beta1.IndexDatapoint.NumericRestriction(
+                                namespace="duration", value_int=10, op="NOT_EQUAL"
+                            ),
                         ],
                     ),
                     per_crowding_attribute_neighbor_count=_TEST_PER_CROWDING_ATTRIBUTE_NUM_NEIGHBOURS,