diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md index db9e1e09e366..52a5ecfe82bd 100644 --- a/sdk/cosmos/azure-cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md @@ -10,6 +10,10 @@ #### Bugs Fixed +* Fixed hybrid-search reciprocal rank fusion (RRF) results being ranked incorrectly when component scores are tied. + Each result now receives a unique ordinal rank, with deterministic tie-breaking based on the decoded document ID + portion of `_rid`. See [PR 49623](https://github.com/Azure/azure-sdk-for-python/pull/49623). + #### Other Changes ### 4.17.1 (2026-09-16) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/aio/hybrid_search_aggregator.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/aio/hybrid_search_aggregator.py index e1a1393238cb..fd24311383c0 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/aio/hybrid_search_aggregator.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/aio/hybrid_search_aggregator.py @@ -5,8 +5,8 @@ """ from azure.cosmos._execution_context.aio.base_execution_context import _QueryExecutionContextBase from azure.cosmos._execution_context.aio import document_producer -from azure.cosmos._execution_context.hybrid_search_aggregator import _retrieve_component_scores, _rewrite_query_infos, \ - _compute_rrf_scores, _compute_ranks, _coalesce_duplicate_rids, _attach_parameters, \ +from azure.cosmos._execution_context.hybrid_search_aggregator import _rewrite_query_infos, \ + _rank_results, _coalesce_duplicate_rids, _attach_parameters, \ _FULL_TEXT_SCORE_SCOPE_KEY, _FULL_TEXT_SCORE_SCOPE_LOCAL, _FULL_TEXT_SCORE_SCOPE_DEFAULT from azure.cosmos._routing import routing_range from azure.cosmos import exceptions @@ -169,35 +169,7 @@ async def _run_hybrid_search(self): # pylint: disable=too-many-branches, too-ma self._format_final_results(drained_results) return - # Get the Components weights if any - if self._hybrid_search_query_info.get('componentWeights'): - component_weights = self._hybrid_search_query_info['componentWeights'] - else: - # If no weights are provided, we default to 1.0 for all components - component_weights = [1.0] * len(self._hybrid_search_query_info['componentQueryInfos']) - - # Sort drained results by _rid - drained_results.sort(key=lambda x: x['_rid']) - - # Compose component scores matrix, where each tuple is (score, index) - component_scores = _retrieve_component_scores(drained_results) - - # Sort by scores using component weights - for index, score_tuples in enumerate(component_scores): - # Negative Weights will change sorting from Descending to Ascending - ordering = self._hybrid_search_query_info['componentQueryInfos'][index]['orderBy'][0] - comparison_factor = not ordering.lower() == 'ascending' - # pylint: disable=cell-var-from-loop - score_tuples.sort(key=lambda x: x[0], reverse=comparison_factor) - - # Compute the ranks - ranks = _compute_ranks(component_scores) - - # Compute the RRF scores and add them to output - _compute_rrf_scores(ranks, component_weights, drained_results) - - # Finally, sort on the RRF scores to build the final result to return - drained_results.sort(key=lambda x: x['Score'], reverse=True) + _rank_results(drained_results, self._hybrid_search_query_info) self._format_final_results(drained_results) def _attach_parameters(self, query): diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/hybrid_search_aggregator.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/hybrid_search_aggregator.py index f85d3eeb60bf..726f3eebde75 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/hybrid_search_aggregator.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/hybrid_search_aggregator.py @@ -3,6 +3,7 @@ """Internal class for multi execution context aggregator implementation in the Azure Cosmos database service. """ +import base64 from typing import Union from azure.cosmos._execution_context.base_execution_context import _QueryExecutionContextBase from azure.cosmos._execution_context import document_producer @@ -24,11 +25,9 @@ class _Placeholders: formattable_order_by = "{documentdb-formattableorderbyquery-filter}" -def _retrieve_component_scores(drained_results): - component_scores_list = [] - for _ in drained_results[0]['payload']['componentScores']: - component_scores_list.append([]) - undefined_components = [-999999] * len(component_scores_list) +def _retrieve_component_scores(drained_results, component_count): + component_scores_list = [[] for _ in range(component_count)] + undefined_components = [-999999] * component_count for index, result in enumerate(drained_results): component_scores = result['payload']['componentScores'] # Another small fix while backend changes are released to deal with empty component score scenarios @@ -56,16 +55,51 @@ def _compute_ranks(component_scores): ranks = [[0] * len(component_scores[0]) for _ in range(len(component_scores))] for component_index, scores in enumerate(component_scores): - rank = 1 # ranks are 1-based for index, score_tuple in enumerate(scores): - # Identical scores should have the same rank - if index > 0 and score_tuple[0] != scores[index - 1][0]: - rank += 1 - ranks[component_index][score_tuple[1]] = rank + ranks[component_index][score_tuple[1]] = index + 1 return ranks +def _get_document_rid(rid): + # Cosmos RID base64 uses '-' for '/', not the base64url alphabet. + decoded = base64.b64decode(rid.replace("-", "/"), validate=True) + if len(decoded) != 16: + raise ValueError("Hybrid search requires a 16-byte document resource ID.") + return int.from_bytes(decoded[8:16], byteorder="little", signed=False) + + +def _rank_results(drained_results, hybrid_search_query_info): + # Use the query-plan weights, defaulting to equal contributions from all components. + component_query_infos = hybrid_search_query_info['componentQueryInfos'] + component_weights = hybrid_search_query_info.get('componentWeights') or [1.0] * len(component_query_infos) + + # Parse once per result, retaining the full RID for identity and deduplication. + document_rids = {result['_rid']: _get_document_rid(result['_rid']) for result in drained_results} + + # Size components from the plan because the first result may have no scores. + # Build (score, result index) lists so component ranks map back to the original results. + component_scores = _retrieve_component_scores(drained_results, len(component_query_infos)) + + # Sort each component by score, breaking ties by numeric RID in the same direction. + for index, score_tuples in enumerate(component_scores): + # Query-plan directions already account for normalized negative weights. + ordering = component_query_infos[index].get('orderBy') or ['Descending'] + score_tuples.sort( + key=lambda score: (score[0], document_rids[drained_results[score[1]]['_rid']]), + reverse=ordering[0].lower() != 'ascending', + ) + + # Assign consecutive ordinal ranks, including for tied component scores. + ranks = _compute_ranks(component_scores) + + # Combine weighted component contributions into each result's RRF score. + _compute_rrf_scores(ranks, component_weights, drained_results) + + # Return highest fused scores first, breaking final ties by numeric RID descending. + drained_results.sort(key=lambda result: (result['Score'], document_rids[result['_rid']]), reverse=True) + + def _coalesce_duplicate_rids(query_results): unique_rids = {d['_rid']: d for d in query_results} return list(unique_rids.values()) @@ -305,36 +339,7 @@ def _run_hybrid_search(self): # pylint: disable=too-many-branches, too-many-sta self._format_final_results(drained_results) return - # Get the Components weight if any - if self._hybrid_search_query_info.get('componentWeights'): - component_weights = self._hybrid_search_query_info['componentWeights'] - else: - # If no weights are provided, we assume all components have equal weight - component_weights = [1.0] * len(self._hybrid_search_query_info['componentQueryInfos']) - - # Sort drained results by _rid - drained_results.sort(key=lambda x: x['_rid']) - - # Compose component scores matrix, where each tuple is (score, index) - component_scores = _retrieve_component_scores(drained_results) - - # Sort by scores using component weights - for index, score_tuples in enumerate(component_scores): - # Ordering of the component query is based on if the weight is negative or positive - # A positive weight ordering means descending order, a negative weight ordering means ascending order - ordering = self._hybrid_search_query_info['componentQueryInfos'][index]['orderBy'][0] - comparison_factor = not ordering.lower() == 'ascending' - # pylint: disable=cell-var-from-loop - score_tuples.sort(key=lambda x: x[0], reverse=comparison_factor) - - # Compute the ranks - ranks = _compute_ranks(component_scores) - - # Compute the RRF scores and add them to output - _compute_rrf_scores(ranks, component_weights, drained_results) - - # Finally, sort on the RRF scores to build the final result to return - drained_results.sort(key=lambda x: x['Score'], reverse=True) + _rank_results(drained_results, self._hybrid_search_query_info) self._format_final_results(drained_results) def _attach_parameters(self, query): diff --git a/sdk/cosmos/azure-cosmos/tests/test_hybrid_search_aggregator.py b/sdk/cosmos/azure-cosmos/tests/test_hybrid_search_aggregator.py index ec7a11fd2713..4746f6c2d136 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_hybrid_search_aggregator.py +++ b/sdk/cosmos/azure-cosmos/tests/test_hybrid_search_aggregator.py @@ -3,17 +3,37 @@ """Deterministic unit tests for client-side hybrid-search ranking.""" +import base64 +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + import pytest +from azure.cosmos.aio import CosmosClient # noqa: F401 - needed to resolve circular imports +from azure.cosmos._execution_context import hybrid_search_aggregator as sync_aggregator +from azure.cosmos._execution_context.aio import hybrid_search_aggregator as async_aggregator from azure.cosmos._execution_context.hybrid_search_aggregator import ( _compute_ranks, _compute_rrf_scores, + _get_document_rid, + _rank_results, ) +pytestmark = pytest.mark.cosmosEmulator + + +def _result(document_rid, scores, payload=None, prefix=b"\x01" * 8): + rid = base64.b64encode(prefix + document_rid.to_bytes(8, "little")).decode("ascii").replace("/", "-") + return {"_rid": rid, "payload": {"componentScores": scores, "payload": payload}} + + +def test_tied_component_scores_receive_consecutive_ordinal_ranks(): + scores = [[(score, index) for index, score in enumerate([10, 5, 5, 5, 0])]] + assert _compute_ranks(scores) == [[1, 2, 3, 4, 5]] + -@pytest.mark.cosmosEmulator def test_mixed_component_rrf_scores_have_exact_expected_order(): - """Pin dense ranking, tied component scores, weighted RRF math, and the + """Pin ordinal ranking, tied component scores, weighted RRF math, and the final descending order without depending on service-generated scores.""" # Each tuple contains (component score, original result index). The first # component is already sorted from highest to lowest, while the second is @@ -28,17 +48,165 @@ def test_mixed_component_rrf_scores_have_exact_expected_order(): _compute_rrf_scores(ranks, [1, 2], results) assert ranks == [ - [1, 2, 2, 3], + [1, 2, 3, 4], [2, 4, 1, 3], ] expected_scores = { "a": 1 / 61 + 2 / 62, "b": 1 / 62 + 2 / 64, - "c": 1 / 62 + 2 / 61, - "d": 1 / 63 + 2 / 63, + "c": 1 / 63 + 2 / 61, + "d": 1 / 64 + 2 / 63, } for result in results: assert result["Score"] == pytest.approx(expected_scores[result["id"]]) results.sort(key=lambda result: result["Score"], reverse=True) - assert [result["id"] for result in results] == ["c", "a", "d", "b"] + assert [result["id"] for result in results] == ["c", "a", "b", "d"] + + +def test_document_rid_is_little_endian_document_portion_not_encoded_order(): + low = _result(255, [], prefix=b"\xff" * 8)["_rid"] + high = _result(256, [], prefix=b"\x00" * 8)["_rid"] + assert "-" in low + assert _get_document_rid(low) == 255 + assert _get_document_rid(high) == 256 + assert _get_document_rid(_result(0x0FEDCBA987654321, [])["_rid"]) == 0x0FEDCBA987654321 + # The prefix is not part of the document number. + assert _get_document_rid(_result(255, [])["_rid"]) == 255 + same_prefix_low = _result(255, [])["_rid"] + same_prefix_high = _result(256, [])["_rid"] + assert same_prefix_low > same_prefix_high + with pytest.raises(ValueError, match="16-byte"): + _get_document_rid(base64.b64encode(b"\x00" * 8).decode("ascii")) + with pytest.raises(ValueError): + _get_document_rid("not a resource ID") + + +@pytest.mark.parametrize("weights", [[1, 2], [1, -2]]) +def test_component_ties_follow_query_plan_direction_and_preserve_weights(weights, monkeypatch): + results = [_result(255, [5, 0.2]), _result(1, [10, 0.1]), _result(256, [5, 0.2])] + info = { + "componentQueryInfos": [{"orderBy": ["Descending"]}, {"orderBy": ["Ascending"]}], + "componentWeights": weights, + } + decode = Mock(wraps=_get_document_rid) + monkeypatch.setattr(sync_aggregator, "_get_document_rid", decode) + _rank_results(results, info) + expected = { + 1: weights[0] / 61 + weights[1] / 61, + 255: weights[0] / 63 + weights[1] / 62, + 256: weights[0] / 62 + weights[1] / 63, + } + assert decode.call_count == len(results) + for result in results: + assert result["Score"] == pytest.approx(expected[_get_document_rid(result["_rid"])]) + assert [_get_document_rid(result["_rid"]) for result in results] == ( + [1, 255, 256] if weights[1] > 0 else [256, 255, 1] + ) + + +def test_final_rrf_ties_are_numeric_rid_descending(): + info = {"componentQueryInfos": [{}, {"orderBy": ["Ascending"]}]} + for incoming in ([1, 255, 256], [256, 255, 1]): + results = [_result(rid, [5, 5]) for rid in incoming] + _rank_results(results, info) + assert [_get_document_rid(result["_rid"]) for result in results] == [256, 1, 255] + scores = {_get_document_rid(result["_rid"]): result["Score"] for result in results} + assert scores[1] == scores[256] + + +def test_zero_score_documents_can_outrank_matches_in_unfiltered_rrf(): + results = [_result(1, [10, 0]), _result(2, [0, 10])] + results.extend(_result(rid, [0, 0]) for rid in (3, 4, 5)) + + _rank_results(results, {"componentQueryInfos": [{}, {}]}) + + # RID 5 ranks second in both components, ahead of documents that rank first + # in one component but last in the other. + assert [_get_document_rid(result["_rid"]) for result in results] == [5, 2, 1, 4, 3] + assert results[0]["Score"] == pytest.approx(2 / 62) + + +@pytest.mark.parametrize("module", [sync_aggregator, async_aggregator], ids=["sync", "async"]) +@pytest.mark.parametrize("legacy_env", [None, "true"], ids=["unset", "legacy-env-ignored"]) +@pytest.mark.asyncio +async def test_query_uses_ordinal_ranks_and_preserves_dedup_payload_and_window(module, legacy_env, monkeypatch): + is_async = module is async_aggregator + if legacy_env is None: + monkeypatch.delenv("AZURE_COSMOS_HYBRID_SEARCH_USE_COMPETITION_RANKING", raising=False) + else: + monkeypatch.setenv("AZURE_COSMOS_HYBRID_SEARCH_USE_COMPETITION_RANKING", legacy_env) + info = { + "requiresGlobalStatistics": False, + "componentQueryInfos": [ + {"orderBy": ["Descending"], "rewrittenQuery": "first"}, + {"orderBy": ["Ascending"], "rewrittenQuery": "second"}, + ], + "skip": 1, + "take": 1, + } + original = _result(1, [5, 5], {"id": "old"}) + replacement = _result(1, [5, 5], {"id": "new"}) + middle = _result(255, [5, 5], "projected scalar") + high = _result(256, [5, 5], {"id": "high"}) + producers = [] + for rows in ([original, middle, high], [replacement, middle]): + peek = AsyncMock(return_value=rows[0]) if is_async else Mock(return_value=rows[0]) + producers.append(SimpleNamespace(peek=peek, _ex_context=SimpleNamespace(_buffer=rows[1:]))) + monkeypatch.setattr(module.document_producer, "_DocumentProducer", Mock(side_effect=producers)) + client = SimpleNamespace(_routing_map_provider=None) + context = module._HybridSearchContextAggregator(client, "colls/test", {}, None, info, None, None) + context._get_target_partition_key_range = AsyncMock(return_value=[{}]) if is_async else Mock(return_value=[{}]) + if is_async: + await context._run_hybrid_search() + actual = [item async for item in context] + else: + context._run_hybrid_search() + actual = list(context) + assert actual == [{"id": "new"}] + assert "Score" not in original + assert "Score" in replacement + assert replacement["_rid"] == original["_rid"] + + +@pytest.mark.parametrize("module", [sync_aggregator, async_aggregator], ids=["sync", "async"]) +@pytest.mark.parametrize("second_scores", [[5, 0.2], []], ids=["empty-first", "all-empty"]) +@pytest.mark.asyncio +async def test_query_uses_plan_component_count_for_empty_scores(module, second_scores, monkeypatch): + is_async = module is async_aggregator + info = { + "requiresGlobalStatistics": False, + "componentQueryInfos": [ + {"orderBy": ["Descending"], "rewrittenQuery": "first"}, + {"orderBy": ["Ascending"], "rewrittenQuery": "second"}, + ], + "componentWeights": [1, 2], + "skip": 0, + "take": 2, + } + empty = _result(255, [], {"id": "empty"}) + second = _result(256, second_scores, {"id": "second"}) + producers = [ + SimpleNamespace( + peek=AsyncMock(return_value=row) if is_async else Mock(return_value=row), + _ex_context=SimpleNamespace(_buffer=[]), + ) + for row in (empty, second) + ] + monkeypatch.setattr(module.document_producer, "_DocumentProducer", Mock(side_effect=producers)) + client = SimpleNamespace(_routing_map_provider=None) + context = module._HybridSearchContextAggregator(client, "colls/test", {}, None, info, None, None) + context._get_target_partition_key_range = AsyncMock(return_value=[{}]) if is_async else Mock(return_value=[{}]) + if is_async: + await context._run_hybrid_search() + actual = [item async for item in context] + else: + context._run_hybrid_search() + actual = list(context) + # The existing -999999 fallback sorts last descending and first ascending; + # when all scores are empty, numeric RID ties give the same ordinal ranks. + assert empty["Score"] == pytest.approx(1 / 62 + 2 / 61) + assert second["Score"] == pytest.approx(1 / 61 + 2 / 62) + assert actual == [{"id": "empty"}, {"id": "second"}] + assert empty["payload"]["componentScores"] == [] + assert second["payload"]["componentScores"] == second_scores diff --git a/sdk/cosmos/azure-cosmos/tests/test_query_hybrid_search.py b/sdk/cosmos/azure-cosmos/tests/test_query_hybrid_search.py index e96bc4405d97..cb8ec95c275b 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_query_hybrid_search.py +++ b/sdk/cosmos/azure-cosmos/tests/test_query_hybrid_search.py @@ -121,6 +121,7 @@ def test_hybrid_search_env_variables(self): os.environ["AZURE_COSMOS_HYBRID_SEARCH_MAX_ITEMS"] = "1000" def test_hybrid_search_queries(self): + expected_rrf_indices = {2, 22, 24, 49, 51, 54, 57, 61, 75, 76, 77, 80, 85} query = "SELECT TOP 10 c.index, c.title FROM c WHERE FullTextContains(c.title, 'John') OR " \ "FullTextContains(c.text, 'John') ORDER BY RANK FullTextScore(c.title, 'John')" results = self.test_container.query_items(query, enable_cross_partition_query=True) @@ -151,13 +152,10 @@ def test_hybrid_search_queries(self): results = self.test_container.query_items(query, enable_cross_partition_query=True) result_list = list(results) assert len(result_list) == 13 - # Captured so the OFFSET query below can be validated against the ranking this same - # container produced, instead of against a hard coded order. Equal RRF scores are broken by - # _rid, which is assigned by the service and differs between containers, so the absolute - # order is not stable across runs - but the offset window must always be a slice of it. + # Tied component scores receive distinct ordinal ranks based on service-assigned RIDs, + # so both the order and top-ten membership can differ between containers. full_rrf_ranking = [res['index'] for res in result_list] - for res in result_list: - assert res['index'] in [61, 51, 49, 54, 75, 24, 77, 76, 80, 25, 22, 2, 66, 57, 85] + assert set(full_rrf_ranking) == expected_rrf_indices query = "SELECT TOP 10 c.index, c.title, c.text FROM c WHERE " \ "FullTextContains(c.title, 'John') OR FullTextContains(c.text, 'John') OR " \ @@ -166,8 +164,9 @@ def test_hybrid_search_queries(self): results = self.test_container.query_items(query, enable_cross_partition_query=True) result_list = list(results) assert len(result_list) == 10 - for res in result_list: - assert res['index'] in [61, 51, 49, 54, 75, 24, 77, 76, 80, 25, 2] + result_indices = {res['index'] for res in result_list} + assert len(result_indices) == 10 + assert result_indices <= expected_rrf_indices query = "SELECT c.index, c.title FROM c WHERE FullTextContains(c.title, 'John')" \ " OR FullTextContains(c.text, 'John') OR FullTextContains(c.text, 'United States') ORDER BY " \ @@ -180,13 +179,21 @@ def test_hybrid_search_queries(self): # RRF ordering regresses, while tolerating service side score/_rid differences. assert [res['index'] for res in result_list] == full_rrf_ranking[5:15] + # Without a WHERE clause, zero-score documents also receive ordinal ranks. + expected_unfiltered_indices = { + item['index'] for item in hybrid_search_data.get_hybrid_search_items()['items'] + } query = "SELECT TOP 10 c.index, c.title FROM c " \ "ORDER BY RANK RRF(FullTextScore(c.title, 'John'), FullTextScore(c.text, 'United States'))" results = self.test_container.query_items(query, enable_cross_partition_query=True) result_list = list(results) assert len(result_list) == 10 - for res in result_list: - assert res['index'] in [61, 51, 49, 54, 75, 24, 77, 76, 80, 25, 2] + result_indices = {res['index'] for res in result_list} + assert len(result_indices) == 10 + assert result_indices <= expected_unfiltered_indices + paged_results = self.test_container.query_items( + query, enable_cross_partition_query=True, max_item_count=3) + assert [res['index'] for res in paged_results] == [res['index'] for res in result_list] query = "SELECT c.index, c.title FROM c " \ "ORDER BY RANK RRF(FullTextScore(c.title, 'John'), FullTextScore(c.text, 'United States')) " \ @@ -194,8 +201,9 @@ def test_hybrid_search_queries(self): results = self.test_container.query_items(query, enable_cross_partition_query=True) result_list = list(results) assert len(result_list) == 13 - for res in result_list: - assert res['index'] in [61, 49, 51, 24, 54, 75, 77, 76, 2, 80, 22, 57, 85] + result_indices = {res['index'] for res in result_list} + assert len(result_indices) == 13 + assert result_indices <= expected_unfiltered_indices item_vector = self.test_container.read_item('50', '1')['vector'] query = "SELECT c.index, c.title FROM c " \ @@ -204,8 +212,9 @@ def test_hybrid_search_queries(self): results = self.test_container.query_items(query, enable_cross_partition_query=True) result_list = list(results) assert len(result_list) == 10 - for res in result_list: - assert res['index'] in [51, 54, 28, 70, 24, 61, 56, 26, 58, 77, 2, 68] + result_indices = {res['index'] for res in result_list} + assert len(result_indices) == 10 + assert result_indices <= expected_unfiltered_indices def test_hybrid_search_query_pagination(self): query = "SELECT c.index, c.title FROM c " \ @@ -244,6 +253,7 @@ def test_hybrid_search_partitioned_query_response_hook(self): assert response_hook.count == 1 def test_hybrid_search_weighted_reciprocal_rank_fusion(self): + expected_rrf_indices = {2, 22, 24, 49, 51, 54, 57, 61, 75, 76, 77, 80, 85} # Test case 1 query = """ SELECT TOP 15 c.index AS Index, c.title AS Title, c.text AS Text @@ -253,9 +263,8 @@ def test_hybrid_search_weighted_reciprocal_rank_fusion(self): """ results = self.test_container.query_items(query, enable_cross_partition_query=True) result_list = [res['Index'] for res in results] - # If some scores rank the same the order of the results may change - for result in result_list: - assert result in [61, 51, 49, 54, 75, 24, 77, 76, 80, 25, 22, 2, 66, 57, 85] + assert len(result_list) == 13 + assert set(result_list) == expected_rrf_indices equal_weights_ranking = result_list # Test case 2 @@ -267,9 +276,6 @@ def test_hybrid_search_weighted_reciprocal_rank_fusion(self): """ results = self.test_container.query_items(query, enable_cross_partition_query=True) result_list = [res['Index'] for res in results] - # If some scores rank the same the order of the results may change - for result in result_list: - assert result in [61, 51, 49, 54, 75, 24, 77, 76, 80, 25, 22, 2, 66, 57, 85] # [10, 10] is a positive uniform scaling of [1, 1], so every RRF score is multiplied by the # same constant and the ranking must be identical. This is an SDK invariant that holds # regardless of the scores the service returns, so it catches weight handling regressions @@ -285,8 +291,13 @@ def test_hybrid_search_weighted_reciprocal_rank_fusion(self): """ results = self.test_container.query_items(query, enable_cross_partition_query=True) result_list = [res['Index'] for res in results] - for result in result_list: - assert result in [61, 51, 49, 54, 75, 24, 77, 76, 80, 2, 25] + assert len(result_list) == 10 + assert len(set(result_list)) == 10 + assert set(result_list) <= expected_rrf_indices + # Keep TOP identical so both queries use the same component candidate limits. + reference_query = query.replace("[0.1, 0.1]", "[1, 1]") + reference_results = self.test_container.query_items(reference_query, enable_cross_partition_query=True) + assert result_list == [res['Index'] for res in reference_results] # Test case 4 query = """ @@ -297,9 +308,8 @@ def test_hybrid_search_weighted_reciprocal_rank_fusion(self): """ results = self.test_container.query_items(query, enable_cross_partition_query=True) result_list = [res['Index'] for res in results] - # If some scores rank the same the order of the results may change - for result in result_list: - assert result in [85, 57, 66, 2, 22, 25, 77, 76, 80, 75, 24, 49, 54, 51, 81, 61] + assert len(result_list) == 13 + assert set(result_list) == expected_rrf_indices # Test case 5 item_vector = self.test_container.read_item('50', '1')['vector'] diff --git a/sdk/cosmos/azure-cosmos/tests/test_query_hybrid_search_async.py b/sdk/cosmos/azure-cosmos/tests/test_query_hybrid_search_async.py index 9d6ff5906cf1..eaec5163d7f3 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_query_hybrid_search_async.py +++ b/sdk/cosmos/azure-cosmos/tests/test_query_hybrid_search_async.py @@ -116,6 +116,7 @@ async def test_hybrid_search_env_variables_async(self): os.environ["AZURE_COSMOS_HYBRID_SEARCH_MAX_ITEMS"] = "1000" async def test_hybrid_search_queries_async(self): + expected_rrf_indices = {2, 22, 24, 49, 51, 54, 57, 61, 75, 76, 77, 80, 85} query = "SELECT TOP 10 c.index, c.title FROM c WHERE FullTextContains(c.title, 'John') OR " \ "FullTextContains(c.text, 'John') ORDER BY RANK FullTextScore(c.title, 'John')" results = self.test_container.query_items(query) @@ -146,13 +147,10 @@ async def test_hybrid_search_queries_async(self): results = self.test_container.query_items(query) result_list = [item async for item in results] assert len(result_list) == 13 - # Captured so the OFFSET query below can be validated against the ranking this same - # container produced, instead of against a hard coded order. Equal RRF scores are broken by - # _rid, which is assigned by the service and differs between containers, so the absolute - # order is not stable across runs - but the offset window must always be a slice of it. + # Tied component scores receive distinct ordinal ranks based on service-assigned RIDs, + # so both the order and top-ten membership can differ between containers. full_rrf_ranking = [res['index'] for res in result_list] - for res in result_list: - assert res['index'] in [61, 51, 49, 54, 75, 24, 77, 76, 80, 25, 22, 2, 66, 57, 85] + assert set(full_rrf_ranking) == expected_rrf_indices query = "SELECT TOP 10 c.index, c.title FROM c WHERE " \ "FullTextContains(c.title, 'John') OR FullTextContains(c.text, 'John') OR " \ @@ -161,8 +159,9 @@ async def test_hybrid_search_queries_async(self): results = self.test_container.query_items(query) result_list = [item async for item in results] assert len(result_list) == 10 - for res in result_list: - assert res['index'] in [61, 51, 49, 54, 75, 24, 77, 76, 80, 25, 2] + result_indices = {res['index'] for res in result_list} + assert len(result_indices) == 10 + assert result_indices <= expected_rrf_indices query = "SELECT c.index, c.title FROM c WHERE FullTextContains(c.title, 'John')" \ " OR FullTextContains(c.text, 'John') OR FullTextContains(c.text, 'United States') ORDER BY " \ @@ -175,13 +174,20 @@ async def test_hybrid_search_queries_async(self): # RRF ordering regresses, while tolerating service side score/_rid differences. assert [res['index'] for res in result_list] == full_rrf_ranking[5:15] + # Without a WHERE clause, zero-score documents also receive ordinal ranks. + expected_unfiltered_indices = { + item['index'] for item in hybrid_search_data.get_hybrid_search_items()['items'] + } query = "SELECT TOP 10 c.index, c.title FROM c " \ "ORDER BY RANK RRF(FullTextScore(c.title, 'John'), FullTextScore(c.text, 'United States'))" results = self.test_container.query_items(query) result_list = [item async for item in results] assert len(result_list) == 10 - for res in result_list: - assert res['index'] in [61, 51, 49, 54, 75, 24, 77, 76, 80, 25, 2] + result_indices = {res['index'] for res in result_list} + assert len(result_indices) == 10 + assert result_indices <= expected_unfiltered_indices + paged_results = self.test_container.query_items(query, max_item_count=3) + assert [res['index'] async for res in paged_results] == [res['index'] for res in result_list] query = "SELECT c.index, c.title FROM c " \ "ORDER BY RANK RRF(FullTextScore(c.title, 'John'), FullTextScore(c.text, 'United States')) " \ @@ -189,8 +195,9 @@ async def test_hybrid_search_queries_async(self): results = self.test_container.query_items(query) result_list = [item async for item in results] assert len(result_list) == 13 - for res in result_list: - assert res['index'] in [61, 49, 51, 24, 54, 75, 77, 76, 2, 80, 22, 57, 85] + result_indices = {res['index'] for res in result_list} + assert len(result_indices) == 13 + assert result_indices <= expected_unfiltered_indices read_item = await self.test_container.read_item('50', '1') item_vector = read_item['vector'] @@ -200,8 +207,9 @@ async def test_hybrid_search_queries_async(self): results = self.test_container.query_items(query) result_list = [item async for item in results] assert len(result_list) == 10 - for res in result_list: - assert res['index'] in [51, 54, 28, 70, 24, 61, 56, 26, 58, 77, 2, 68] + result_indices = {res['index'] for res in result_list} + assert len(result_indices) == 10 + assert result_indices <= expected_unfiltered_indices async def test_hybrid_search_query_pagination_async(self): query = "SELECT c.index, c.title FROM c " \ @@ -242,6 +250,7 @@ async def test_hybrid_search_partitioned_query_response_hook_async(self): assert response_hook.count == 1 async def test_hybrid_search_weighted_reciprocal_rank_fusion_async(self): + expected_rrf_indices = {2, 22, 24, 49, 51, 54, 57, 61, 75, 76, 77, 80, 85} # Test case 1 query = """ SELECT TOP 15 c.index AS Index, c.title AS Title, c.text AS Text @@ -251,9 +260,8 @@ async def test_hybrid_search_weighted_reciprocal_rank_fusion_async(self): """ results = self.test_container.query_items(query) result_list = [res['Index'] async for res in results] - # If some scores rank the same the order of the results may change - for result in result_list: - assert result in [61, 51, 49, 54, 75, 24, 77, 76, 80, 25, 22, 2, 66, 57, 85] + assert len(result_list) == 13 + assert set(result_list) == expected_rrf_indices equal_weights_ranking = result_list # Test case 2 @@ -265,9 +273,6 @@ async def test_hybrid_search_weighted_reciprocal_rank_fusion_async(self): """ results = self.test_container.query_items(query) result_list = [res['Index'] async for res in results] - # If some scores rank the same the order of the results may change - for result in result_list: - assert result in [61, 51, 49, 54, 75, 24, 77, 76, 80, 25, 22, 2, 66, 57, 85] # [10, 10] is a positive uniform scaling of [1, 1], so every RRF score is multiplied by the # same constant and the ranking must be identical. This is an SDK invariant that holds # regardless of the scores the service returns, so it catches weight handling regressions @@ -283,8 +288,13 @@ async def test_hybrid_search_weighted_reciprocal_rank_fusion_async(self): """ results = self.test_container.query_items(query) result_list = [res['Index'] async for res in results] - for result in result_list: - assert result in [61, 51, 49, 54, 75, 24, 77, 76, 80, 2, 25] + assert len(result_list) == 10 + assert len(set(result_list)) == 10 + assert set(result_list) <= expected_rrf_indices + # Keep TOP identical so both queries use the same component candidate limits. + reference_query = query.replace("[0.1, 0.1]", "[1, 1]") + reference_results = self.test_container.query_items(reference_query) + assert result_list == [res['Index'] async for res in reference_results] # Test case 4 query = """ @@ -295,9 +305,8 @@ async def test_hybrid_search_weighted_reciprocal_rank_fusion_async(self): """ results = self.test_container.query_items(query) result_list = [res['Index'] async for res in results] - # If some scores rank the same the order of the results may change - for result in result_list: - assert result in [85, 57, 66, 2, 22, 25, 77, 76, 80, 75, 24, 49, 54, 51, 81, 61] + assert len(result_list) == 13 + assert set(result_list) == expected_rrf_indices # Test case 5 read_item = await self.test_container.read_item('50', '1')