Skip to content

Commit d418ba1

Browse files
vertex-sdk-botcopybara-github
authored andcommitted
fix: Preserve rrf_ranking_alpha=0.0 in MatchingEngineIndexEndpoint hybrid queries
`HybridQuery.rrf_ranking_alpha` is an optional float that defaults to `None`. Two call sites in `matching_engine_index_endpoint.py` previously gated RRF wiring on the truthiness of this field (`if query.rrf_ranking_alpha`), which incorrectly treated an explicit `0.0` the same as an unset value. As a result, callers passing `rrf_ranking_alpha=0.0` (pure sparse ranking) had the `rrf` submessage silently omitted from the outgoing request, and the server fell back to its default alpha instead of honoring the caller's intent. Switch both sites to `is not None` so `0.0` is preserved end-to-end: - `find_neighbors` (public path, `FindNeighborsRequest.Query.RRF`) - `match` (private-service-access path, `MatchRequest.RRF`) Extend `_TEST_HYBRID_QUERIES` with a sparse-only entry that sets `rrf_ranking_alpha=0.0`, and update the two consumers (`test_private_service_access_hybrid_search_match_queries` and `test_index_public_endpoint_find_neighbors_queries`) to assert that the resulting RPC contains `rrf { alpha: 0.0 }` (instead of no rrf at all). Also switch the expected-side predicate in the private-service-access test from truthiness to `is not None` for consistency with the SDK. FUTURE_COPYBARA_INTEGRATE_REVIEW=#6967 from googleapis:release-please--branches--main 4a0d1d2 PiperOrigin-RevId: 948540018
1 parent 3335750 commit d418ba1

2 files changed

Lines changed: 29 additions & 3 deletions

File tree

google/cloud/aiplatform/matching_engine/matching_engine_index_endpoint.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1865,7 +1865,7 @@ def find_neighbors(
18651865
dimensions=query.sparse_embedding_dimensions,
18661866
),
18671867
)
1868-
if query.rrf_ranking_alpha:
1868+
if query.rrf_ranking_alpha is not None:
18691869
find_neighbors_query.rrf = (
18701870
gca_match_service_v1beta1.FindNeighborsRequest.Query.RRF(
18711871
alpha=query.rrf_ranking_alpha,
@@ -2200,7 +2200,7 @@ def match(
22002200
match_service_pb2.MatchRequest.RRF(
22012201
alpha=query.rrf_ranking_alpha,
22022202
)
2203-
if query_is_hybrid and query.rrf_ranking_alpha
2203+
if query_is_hybrid and query.rrf_ranking_alpha is not None
22042204
else None
22052205
),
22062206
)

tests/unit/aiplatform/test_matching_engine_index_endpoint.py

Lines changed: 27 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -258,6 +258,11 @@
258258
sparse_embedding_dimensions=[1, 2, 3],
259259
sparse_embedding_values=[0.1, 0.2, 0.3],
260260
),
261+
HybridQuery(
262+
sparse_embedding_dimensions=[1, 2, 3],
263+
sparse_embedding_values=[0.1, 0.2, 0.3],
264+
rrf_ranking_alpha=0.0, # Test edge case where alpha is 0
265+
),
261266
]
262267
_TEST_NUM_NEIGHBOURS = 1
263268
_TEST_FILTER = [
@@ -1539,7 +1544,7 @@ def test_private_service_access_hybrid_search_match_queries(
15391544
match_service_pb2.MatchRequest.RRF(
15401545
alpha=_TEST_HYBRID_QUERIES[i].rrf_ranking_alpha,
15411546
)
1542-
if _TEST_HYBRID_QUERIES[i].rrf_ranking_alpha
1547+
if _TEST_HYBRID_QUERIES[i].rrf_ranking_alpha is not None
15431548
else None
15441549
),
15451550
)
@@ -2093,6 +2098,27 @@ def test_index_public_endpoint_find_neighbors_queries(
20932098
approximate_neighbor_count=_TEST_APPROX_NUM_NEIGHBORS,
20942099
fraction_leaf_nodes_to_search_override=_TEST_FRACTION_LEAF_NODES_TO_SEARCH_OVERRIDE,
20952100
),
2101+
gca_match_service_v1beta1.FindNeighborsRequest.Query(
2102+
neighbor_count=_TEST_NUM_NEIGHBOURS,
2103+
datapoint=gca_index_v1beta1.IndexDatapoint(
2104+
restricts=[
2105+
gca_index_v1beta1.IndexDatapoint.Restriction(
2106+
namespace="class",
2107+
allow_list=["token_1"],
2108+
deny_list=["token_2"],
2109+
)
2110+
],
2111+
sparse_embedding=gca_index_v1beta1.IndexDatapoint.SparseEmbedding(
2112+
values=[0.1, 0.2, 0.3], dimensions=[1, 2, 3]
2113+
),
2114+
),
2115+
rrf=gca_match_service_v1beta1.FindNeighborsRequest.Query.RRF(
2116+
alpha=0.0,
2117+
),
2118+
per_crowding_attribute_neighbor_count=_TEST_PER_CROWDING_ATTRIBUTE_NUM_NEIGHBOURS,
2119+
approximate_neighbor_count=_TEST_APPROX_NUM_NEIGHBORS,
2120+
fraction_leaf_nodes_to_search_override=_TEST_FRACTION_LEAF_NODES_TO_SEARCH_OVERRIDE,
2121+
),
20962122
],
20972123
return_full_datapoint=_TEST_RETURN_FULL_DATAPOINT,
20982124
)

0 commit comments

Comments
 (0)