From 4eb5fc9bd9041033e74baaa6e29c7988f94a6662 Mon Sep 17 00:00:00 2001 From: Maya Painter Date: Thu, 22 Oct 2026 13:51:00 -0700 Subject: [PATCH 1/5] [Cosmos] Use deterministic ordinal ranks for hybrid-search RRF Port ordinal and opt-in competition ranking with cached numeric document RID tie-breaking and offline regression coverage. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 29c3522f-9887-43c5-bebb-f42ee9ba4d51 --- .../query/HybridSearchRankingTest.java | 295 ++++++++++++++++++ sdk/cosmos/azure-cosmos/CHANGELOG.md | 1 + .../azure/cosmos/implementation/Configs.java | 10 + ...idSearchDocumentQueryExecutionContext.java | 68 ++-- .../hybridsearch/HybridSearchQueryResult.java | 8 + 5 files changed, 359 insertions(+), 23 deletions(-) create mode 100644 sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/query/HybridSearchRankingTest.java diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/query/HybridSearchRankingTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/query/HybridSearchRankingTest.java new file mode 100644 index 0000000000000..7e09592338203 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/query/HybridSearchRankingTest.java @@ -0,0 +1,295 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.query; + +import com.azure.cosmos.implementation.Configs; +import com.azure.cosmos.implementation.Document; +import com.azure.cosmos.implementation.ResourceType; +import com.azure.cosmos.implementation.query.HybridSearchDocumentQueryExecutionContext.ComponentWeight; +import com.azure.cosmos.implementation.query.HybridSearchDocumentQueryExecutionContext.ScoreTuple; +import com.azure.cosmos.implementation.query.hybridsearch.HybridSearchQueryInfo; +import com.azure.cosmos.implementation.query.hybridsearch.HybridSearchQueryResult; +import com.azure.cosmos.models.CosmosQueryRequestOptions; +import com.azure.cosmos.models.FeedResponse; +import com.azure.cosmos.models.SqlQuerySpec; +import org.testng.annotations.DataProvider; +import org.testng.annotations.Test; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; + +import java.lang.reflect.Field; +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Base64; +import java.util.Collections; +import java.util.List; +import java.util.UUID; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; + +import static com.azure.cosmos.implementation.query.HybridSearchDocumentQueryExecutionContext.coalesceAndSortResults; +import static com.azure.cosmos.implementation.query.HybridSearchDocumentQueryExecutionContext.computeRRFScores; +import static com.azure.cosmos.implementation.query.HybridSearchDocumentQueryExecutionContext.computeRanks; +import static com.azure.cosmos.implementation.query.HybridSearchDocumentQueryExecutionContext.retrieveComponentScores; +import static com.azure.cosmos.implementation.query.HybridSearchDocumentQueryExecutionContext.retrieveComponentWeights; +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.within; +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.spy; + +public class HybridSearchRankingTest { + private static final String RANKING_PROPERTY = "COSMOS.HYBRID_SEARCH_USE_COMPETITION_RANKING"; + + @DataProvider(name = "rankingModes") + public Object[][] rankingModes() { + return new Object[][] { + {false, Arrays.asList(1, 2, 3, 4, 5)}, + {true, Arrays.asList(1, 2, 2, 2, 5)} + }; + } + + @Test(groups = "unit", dataProvider = "rankingModes") + public void tiedScoresUseOrdinalOrCompetitionRanks(boolean competition, List expected) { + List> rows = Arrays.asList( + row(5, 10), row(4, 5), row(3, 5), row(2, 5), row(1, 0)); + List weights = weights(null, SortOrder.Descending); + + List> ranks = ranks(rows, weights, competition); + assertThat(ranks).containsExactly(expected); + + List> fused = computeRRFScores( + Mono.just(ranks), Mono.just(new ArrayList<>(rows)), weights).collectList().block(); + assertThat(fused).containsExactlyElementsOf(rows); + for (int i = 0; i < rows.size(); i++) { + assertThat(rows.get(i).getScore()).isCloseTo(1.0 / (60 + expected.get(i)), within(1e-15)); + } + } + + @DataProvider(name = "directions") + public Object[][] directions() { + return new Object[][] { + {SortOrder.Ascending, Arrays.asList(1, 2)}, + {SortOrder.Descending, Arrays.asList(2, 1)} + }; + } + + @Test(groups = "unit", dataProvider = "directions") + public void tiesFollowPlanDirectionRatherThanEncodedRidOrWeight( + SortOrder direction, List expected) { + HybridSearchQueryResult low = row(1, 5); + HybridSearchQueryResult high = row(256, 5); + assertThat(low.getRid().compareTo(high.getRid())).isPositive(); + + // A negative original weight can arrive normalized to positive with an ascending plan. + List weights = weights(Collections.singletonList(2.0), direction); + assertThat(ranks(Arrays.asList(low, high), weights, false)).containsExactly(expected); + assertThat(ranks(Arrays.asList(high, low), weights, false)) + .containsExactly(Arrays.asList(expected.get(1), expected.get(0))); + assertThat(ranks(Arrays.asList(low, high), weights, true)) + .containsExactly(Arrays.asList(1, 1)); + } + + @Test(groups = "unit") + public void mixedDirectionsPreserveWeightsAndFormula() { + List> rows = Arrays.asList(row(1, 10, 5), row(256, 5, 10)); + List weights = weights(Arrays.asList(2.0, 3.0), + SortOrder.Descending, SortOrder.Ascending); + List> ranks = ranks(rows, weights, false); + + assertThat(ranks).containsExactly(Arrays.asList(1, 2), Arrays.asList(1, 2)); + computeRRFScores(Mono.just(ranks), Mono.just(rows), weights).collectList().block(); + assertThat(rows.get(0).getScore()).isCloseTo(5.0 / 61, within(1e-15)); + assertThat(rows.get(1).getScore()).isCloseTo(5.0 / 62, within(1e-15)); + } + + @Test(groups = "unit") + public void finalScoreTiesUseDescendingNumericDocumentRid() { + List> rows = Arrays.asList(row(1, 10, 5), row(256, 5, 10)); + List weights = weights(null, SortOrder.Descending, SortOrder.Descending); + List> fused = computeRRFScores( + Mono.just(ranks(rows, weights, false)), Mono.just(rows), weights).collectList().block(); + + assertThat(fused).extracting(HybridSearchQueryResult::getDocumentRid).containsExactly(256L, 1L); + assertThat(fused.get(0).getScore()).isEqualTo(fused.get(1).getScore()); + } + + @Test(groups = "unit") + public void ridDecodingUsesLittleEndianAndCosmosBase64() { + long documentRid = 0x000102030405ffffL; + HybridSearchQueryResult result = row(documentRid, 5); + + assertThat(result.getRid()).contains("-"); + assertThat(result.getDocumentRid()).isEqualTo(documentRid); + assertThat(row(128, 5).getDocumentRid()).isEqualTo(128); + } + + @Test(groups = "unit") + public void numericRidIsCachedBeforeRankingAndFusion() { + HybridSearchQueryResult low = row(1, 5); + HybridSearchQueryResult high = row(256, 5); + low.rid = "not-a-rid"; + high.rid = "not-a-rid-either"; + List> rows = Arrays.asList(low, high); + List weights = weights(null, SortOrder.Descending); + + assertThat(low.getDocumentRid()).isEqualTo(1); + assertThat(high.getDocumentRid()).isEqualTo(256); + assertThat(ranks(rows, weights, false)).containsExactly(Arrays.asList(2, 1)); + List> fused = computeRRFScores( + Mono.just(ranks(rows, weights, true)), Mono.just(rows), weights).collectList().block(); + assertThat(fused).containsExactly(high, low); + } + + @Test(groups = "unit") + public void numericRidComparisonsAreUnsigned() { + HybridSearchQueryResult high = spy(row(256, 5)); + doReturn(-1L).when(high).getDocumentRid(); + HybridSearchQueryResult low = row(1, 5); + List> rows = Arrays.asList(low, high); + List weights = weights(null, SortOrder.Descending); + + assertThat(ranks(rows, weights, false)).containsExactly(Arrays.asList(2, 1)); + List> fused = computeRRFScores( + Mono.just(ranks(rows, weights, true)), Mono.just(rows), weights).collectList().block(); + assertThat(fused).containsExactly(high, low); + } + + @Test(groups = "unit") + public void deduplicationKeepsEncodedIdentityAndFirstPayload() { + String rid = rid(1, 1); + Document first = document(rid, "first", 10); + Document duplicate = document(rid, "duplicate", 100); + Document otherCollection = document(rid(2, 1), "other-collection", 5); + Document high = document(rid(1, 256), "high", 3); + + List> rows = coalesceAndSortResults( + Flux.just(first, duplicate, otherCollection, high)).block(); + + assertThat(rows).hasSize(3); + assertThat(rows).extracting(HybridSearchQueryResult::getDocumentRid).containsExactly(256L, 1L, 1L); + assertThat(rows).extracting(result -> result.getPayload().getId()) + .containsExactly("high", "first", "other-collection"); + assertThat(rows.get(1).getComponentScores()).containsExactly(10.0); + } + + @Test(groups = "unit", dataProvider = "rankingModes") + public void emptyScoresUsePlanDimensionsAndExistingUndefinedValue( + boolean competition, List ignored) { + List> rows = Arrays.asList(row(256), row(1, 10, 5)); + List weights = weights(null, SortOrder.Descending, SortOrder.Ascending); + List> scores = retrieveComponentScores(Mono.just(rows), weights).block(); + + assertThat(scores).hasSize(2); + assertThat(scores.get(0)).extracting(ScoreTuple::getScore).containsExactly(10.0, -999999.0); + assertThat(scores.get(1)).extracting(ScoreTuple::getScore).containsExactly(-999999.0, 5.0); + assertThat(ranks(rows, weights, competition)).containsExactly(Arrays.asList(2, 1), Arrays.asList(1, 2)); + + Collections.reverse(rows); + assertThat(ranks(rows, weights, competition)).containsExactly(Arrays.asList(1, 2), Arrays.asList(2, 1)); + } + + @Test(groups = "unit") + public void emptyInputAndAllEmptyScoresAreSupported() { + List weights = weights(null, SortOrder.Descending, SortOrder.Descending); + Mono>> empty = coalesceAndSortResults(Flux.empty()).cache(); + Mono>> emptyRanks = computeRanks(retrieveComponentScores(empty, weights), false); + assertThat(computeRRFScores(emptyRanks, empty, weights).collectList().block()).isEmpty(); + + List> rows = Arrays.asList(row(256), row(1)); + assertThat(ranks(rows, weights, false)).containsExactly(Arrays.asList(1, 2), Arrays.asList(1, 2)); + assertThat(ranks(rows, weights, true)).containsExactly(Arrays.asList(1, 1), Arrays.asList(1, 1)); + assertThat(ranks(Collections.singletonList(row(1)), weights, false)) + .containsExactly(Collections.singletonList(1), Collections.singletonList(1)); + } + + @Test(groups = "unit") + public void rankingReusesCollectedRowsAndPreservesPayloadPaging() throws ReflectiveOperationException { + AtomicInteger subscriptions = new AtomicInteger(); + Mono>> rows = coalesceAndSortResults(Flux.defer(() -> { + subscriptions.incrementAndGet(); + return Flux.just(document(rid(1, 1), "low", 5), document(rid(1, 256), "high", 5), + document(rid(1, 128), "middle", 5)); + })).cache(); + List weights = weights(null, SortOrder.Descending); + Flux> fused = computeRRFScores( + computeRanks(retrieveComponentScores(rows, weights), false), rows, weights); + HybridSearchDocumentQueryExecutionContext context = context(); + Field observable = HybridSearchDocumentQueryExecutionContext.class.getDeclaredField("hybridObservable"); + observable.setAccessible(true); + observable.set(context, fused); + + List> pages = context.drainAsync(2).collectList().block(); + + assertThat(subscriptions.get()).isEqualTo(1); + assertThat(pages).hasSize(2); + assertThat(pages.get(0).getResults()).extracting(Document::getId).containsExactly("high", "middle"); + assertThat(pages.get(1).getResults()).extracting(Document::getId).containsExactly("low"); + } + + @Test(groups = "unit") + public void configurationIsResolvedOncePerQuery() throws ReflectiveOperationException { + String previous = System.getProperty(RANKING_PROPERTY); + try { + System.clearProperty(RANKING_PROPERTY); + assertThat(Configs.isHybridSearchCompetitionRankingEnabled()).isEqualTo( + Boolean.parseBoolean(System.getenv("AZURE_COSMOS_HYBRID_SEARCH_USE_COMPETITION_RANKING"))); + + System.setProperty(RANKING_PROPERTY, "false"); + HybridSearchDocumentQueryExecutionContext ordinal = context(); + System.setProperty(RANKING_PROPERTY, "true"); + HybridSearchDocumentQueryExecutionContext competition = context(); + Field rankingFlag = HybridSearchDocumentQueryExecutionContext.class.getDeclaredField("useCompetitionRanking"); + rankingFlag.setAccessible(true); + assertThat(rankingFlag.getBoolean(ordinal)).isFalse(); + assertThat(rankingFlag.getBoolean(competition)).isTrue(); + } finally { + if (previous == null) { + System.clearProperty(RANKING_PROPERTY); + } else { + System.setProperty(RANKING_PROPERTY, previous); + } + } + } + + private static HybridSearchDocumentQueryExecutionContext context() { + return new HybridSearchDocumentQueryExecutionContext(null, mock(IDocumentQueryClient.class), + ResourceType.Document, new SqlQuerySpec("SELECT * FROM c"), new CosmosQueryRequestOptions(), + "dbs/db/colls/coll", null, UUID.randomUUID(), new AtomicBoolean(), new HybridSearchQueryInfo()); + } + + private static List> ranks(List> rows, + List weights, boolean competition) { + return computeRanks(retrieveComponentScores(Mono.just(rows), weights), competition).block(); + } + + private static List weights(List values, SortOrder... directions) { + List components = new ArrayList<>(); + for (SortOrder direction : directions) { + QueryInfo info = new QueryInfo(); + info.setOrderBy(Collections.singletonList(direction)); + components.add(info); + } + return retrieveComponentWeights(values, components); + } + + private static HybridSearchQueryResult row(long documentRid, double... scores) { + return new HybridSearchQueryResult<>(document(rid(1, documentRid), "doc-" + documentRid, scores).toJson()); + } + + private static Document document(String rid, String id, double... scores) { + return new Document("{\"_rid\":\"" + rid + "\",\"payload\":{\"componentScores\":" + + Arrays.toString(scores) + ",\"payload\":{\"id\":\"" + id + "\"}}}"); + } + + private static String rid(int database, long document) { + byte[] bytes = new byte[16]; + ByteBuffer.wrap(bytes).putInt(database); + bytes[4] = (byte) 0x80; + ByteBuffer.wrap(bytes, 8, 8).order(ByteOrder.LITTLE_ENDIAN).putLong(document); + return Base64.getEncoder().encodeToString(bytes).replace('/', '-'); + } +} diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md index 2cab777b4aeef..79ac619c096c7 100644 --- a/sdk/cosmos/azure-cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md @@ -7,6 +7,7 @@ #### Breaking Changes #### Bugs Fixed +* Hybrid search now uses deterministic ordinal reciprocal-rank fusion: component scores `10, 5, 5, 5, 0` receive ranks `1, 2, 3, 4, 5`. Set `AZURE_COSMOS_HYBRID_SEARCH_USE_COMPETITION_RANKING=true` to use competition ranks `1, 2, 2, 2, 5` instead. Component ties use numeric document RID in the query-plan sort direction; fused-score ties use numeric document RID descending. * Fixed the Direct mode (RNTBD) encoding of the `x-ms-cosmos-workload-id` header, which was sent with the wrong token ID (`0x00DC`) and type (`Byte`) instead of `0x00E7` and `UShort`. - See [PR 50673](https://github.com/Azure/azure-sdk-for-java/pull/50673) #### Other Changes diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java index b152a8f6e3b02..babf448ce8b99 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java @@ -269,6 +269,11 @@ public class Configs { private static final String MAX_ITEM_COUNT_FOR_HYBRID_SEARCH_VARIABLE = "COSMOS_MAX_ITEM_SIZE_FOR_HYBRID_SEARCH"; public static final int DEFAULT_MAX_ITEM_COUNT_FOR_HYBRID_SEARCH = 1000; + private static final String HYBRID_SEARCH_USE_COMPETITION_RANKING = + "COSMOS.HYBRID_SEARCH_USE_COMPETITION_RANKING"; + private static final String HYBRID_SEARCH_USE_COMPETITION_RANKING_VARIABLE = + "AZURE_COSMOS_HYBRID_SEARCH_USE_COMPETITION_RANKING"; + private static final String AZURE_COSMOS_DISABLE_NON_STREAMING_ORDER_BY = "COSMOS.AZURE_COSMOS_DISABLE_NON_STREAMING_ORDER_BY"; private static final boolean DEFAULT_AZURE_COSMOS_DISABLE_NON_STREAMING_ORDER_BY = false; @@ -1248,6 +1253,11 @@ public static int getMaxItemCountForHybridSearchSearch() { return DEFAULT_MAX_ITEM_COUNT_FOR_HYBRID_SEARCH; } + public static boolean isHybridSearchCompetitionRankingEnabled() { + return Boolean.parseBoolean(System.getProperty(HYBRID_SEARCH_USE_COMPETITION_RANKING, + firstNonNull(emptyToNull(System.getenv(HYBRID_SEARCH_USE_COMPETITION_RANKING_VARIABLE)), "false"))); + } + public static boolean getAzureCosmosNonStreamingOrderByDisabled() { if(logger.isTraceEnabled()) { logger.trace( diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/HybridSearchDocumentQueryExecutionContext.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/HybridSearchDocumentQueryExecutionContext.java index b05249db89ec3..3431df1f9a542 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/HybridSearchDocumentQueryExecutionContext.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/HybridSearchDocumentQueryExecutionContext.java @@ -7,6 +7,7 @@ import com.azure.cosmos.CosmosDiagnostics; import com.azure.cosmos.CosmosException; import com.azure.cosmos.implementation.ClientSideRequestStatistics; +import com.azure.cosmos.implementation.Configs; import com.azure.cosmos.implementation.DiagnosticsClientContext; import com.azure.cosmos.implementation.Document; import com.azure.cosmos.implementation.DocumentClientRetryPolicy; @@ -78,6 +79,7 @@ private static ImplementationBridgeHelpers.FeedResponseHelper.FeedResponseAccess private Flux> hybridObservable; private Mono aggregatedGlobalStatistics; private final SchedulingStopwatch hybridSearchSchedulingStopwatch; + private final boolean useCompetitionRanking; protected HybridSearchDocumentQueryExecutionContext( DiagnosticsClientContext diagnosticsClientContext, @@ -97,6 +99,7 @@ protected HybridSearchDocumentQueryExecutionContext( this.tracker = new RequestChargeTracker(); this.queryMetricMap = new ConcurrentHashMap<>(); this.clientSideRequestStatistics = ConcurrentHashMap.newKeySet(); + this.useCompetitionRanking = Configs.isHybridSearchCompetitionRankingEnabled(); // Initialize the shared stopwatch for hybrid search timing this.hybridSearchSchedulingStopwatch = new HybridSearchCumulativeSchedulingStopWatch(); @@ -196,15 +199,15 @@ private Flux> hybridSearch(List componentQueryResults = getComponentQueryResults(targetFeedRanges, initialPageSize, collection, rewrittenQueryInfos); - // Coalesce the results on unique _rids, and sort it based on the _rid + // Deduplicate by encoded RID, then order by numeric document RID. Mono>> coalescedAndSortedResults = coalesceAndSortResults(componentQueryResults).cache(); // Compose component scores matrix, where each tuple is (score, index) Mono>> componentScoresList = retrieveComponentScores(coalescedAndSortedResults, componentWeights); - // Compute Ranks - Mono>> ranks = computeRanks(componentScoresList); + // Assign ordinal ranks by default, or competition ranks for equal component scores. + Mono>> ranks = computeRanks(componentScoresList, this.useCompetitionRanking); // Compute the RRF scores return computeRRFScores(ranks, coalescedAndSortedResults, componentWeights) @@ -332,9 +335,10 @@ public Flux> apply(Flux } } - private static Flux> computeRRFScores(Mono>> ranks, - Mono>> coalescedAndSortedResults, - List componentWeights) { + static Flux> computeRRFScores( + Mono>> ranks, + Mono>> coalescedAndSortedResults, + List componentWeights) { return ranks.zipWith(coalescedAndSortedResults) .map(tuple -> { List> ranksInternal = tuple.getT1(); @@ -347,13 +351,19 @@ private static Flux> computeRRFScores(Mono { + int scoreComparison = Double.compare(y.getScore(), x.getScore()); + return scoreComparison != 0 ? scoreComparison + : Long.compareUnsigned(y.getDocumentRid(), x.getDocumentRid()); + }); return results; }).flatMapMany(Flux::fromIterable); } - private static Mono>> computeRanks(Mono>> componentScoresList) { + static Mono>> computeRanks( + Mono>> componentScoresList, boolean useCompetitionRanking) { return componentScoresList.map(componentScores -> { // initialize ranks as an N-D list with zeros List> ranksInternal = new ArrayList<>(); @@ -368,9 +378,11 @@ private static Mono>> computeRanks(Mono for (int componentIndex = 0; componentIndex < componentScores.size(); componentIndex++) { int rank = 1; // ranks are 1 based for (int index = 0; index < componentScores.get(componentIndex).size(); index++) { - // Identical scores should have the same rank - if ((index > 0) && (componentScores.get(componentIndex).get(index).getScore() != componentScores.get(componentIndex).get(index - 1).getScore())) { - rank += 1; + // Competition ties share the first position; the next distinct score skips tied positions. + if (!useCompetitionRanking || (index > 0 + && componentScores.get(componentIndex).get(index).getScore().doubleValue() + != componentScores.get(componentIndex).get(index - 1).getScore().doubleValue())) { + rank = index + 1; } int rankIndex = componentScores.get(componentIndex).get(index).getIndex(); ranksInternal.get(componentIndex).set(rankIndex, rank); @@ -380,10 +392,12 @@ private static Mono>> computeRanks(Mono }); } - private static Mono>> retrieveComponentScores(Mono>> coalescedAndSortedResults, List componentWeights) { + static Mono>> retrieveComponentScores( + Mono>> coalescedAndSortedResults, List componentWeights) { return coalescedAndSortedResults.map(results -> { List> componentScoresInternal = new ArrayList<>(); - for (int i = 0; i < results.get(0).getComponentScores().size(); i++) { + // The plan defines the dimensions, even when the first row has no component scores. + for (int i = 0; i < componentWeights.size(); i++) { componentScoresInternal.add(new ArrayList<>()); } List undefinedComponentScores = new ArrayList<>(); @@ -401,10 +415,16 @@ private static Mono>> retrieveComponentScores(Mono - componentWeights.get(componentIndex).getComparator().compare(x.getScore(), y.getScore())); + ComponentWeight componentWeight = componentWeights.get(i); + componentScoresInternal.get(i).sort((x, y) -> { + int scoreComparison = componentWeight.getComparator().compare(x.getScore(), y.getScore()); + return scoreComparison != 0 ? scoreComparison + : componentWeight.comparisonFactor * Long.compareUnsigned( + results.get(x.getIndex()).getDocumentRid(), results.get(y.getIndex()).getDocumentRid()); + }); } return componentScoresInternal; }); @@ -499,7 +519,7 @@ private Flux retrieveRewrittenQueryInfos(List componentQue }); } - private Mono>> coalesceAndSortResults(Flux componentQueryResults) { + static Mono>> coalesceAndSortResults(Flux componentQueryResults) { return componentQueryResults.collectList() .map(results -> { Map> uniqueDocuments = new LinkedHashMap<>(); @@ -509,9 +529,9 @@ private Mono>> coalesceAndSortResults(Flu uniqueDocuments.putIfAbsent(rid, result); } - // Sort by _rid + // RID identity remains encoded; only ordering uses the cached document component. List> coalescedResults = new ArrayList<>(uniqueDocuments.values()); - coalescedResults.sort(Comparator.comparing(HybridSearchQueryResult::getRid)); + coalescedResults.sort((x, y) -> Long.compareUnsigned(y.getDocumentRid(), x.getDocumentRid())); return coalescedResults; }); } @@ -571,7 +591,7 @@ private GlobalFullTextSearchQueryStatistics aggregateStatistics(List retrieveComponentWeights(List componentWeightList, List componentQueryInfos) { + static List retrieveComponentWeights(List componentWeightList, List componentQueryInfos) { boolean useDefaultComponentWeight = componentWeightList == null || componentWeightList.isEmpty(); List componentWeights = new ArrayList<>(); for (int i=0;i retrieveComponentWeights(List componentWei return componentWeights; } - private static class ComponentWeight { + static class ComponentWeight { private final Double weight; private final Comparator comparator; + private final int comparisonFactor; public ComponentWeight(Double weight, SortOrder sortOrder) { this.weight = weight; - int comparisonFactor = (sortOrder == SortOrder.Ascending) ? 1 : -1; + // Direction comes from the plan, not the (potentially normalized) weight. + this.comparisonFactor = (sortOrder == SortOrder.Ascending) ? 1 : -1; this.comparator = (x, y) -> comparisonFactor * Double.compare(x, y); } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/hybridsearch/HybridSearchQueryResult.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/hybridsearch/HybridSearchQueryResult.java index 8634e3b2cfd6b..0b2b088145777 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/hybridsearch/HybridSearchQueryResult.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/hybridsearch/HybridSearchQueryResult.java @@ -5,6 +5,7 @@ import com.azure.cosmos.implementation.Constants; import com.azure.cosmos.implementation.Document; +import com.azure.cosmos.implementation.ResourceId; import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.node.ObjectNode; @@ -16,11 +17,14 @@ public class HybridSearchQueryResult extends Document { public List componentScores; public Document payload; public Double score; + private final long documentRid; public HybridSearchQueryResult(String jsonString) { super(jsonString); // Initialize rid this.rid = super.getResourceId(); + // getDocument exposes the little-endian document component; compare it as unsigned. + this.documentRid = ResourceId.parse(this.rid).getDocument(); // Initialize componentScores final Object outerObject = super.get(Constants.Properties.PAYLOAD); if (outerObject instanceof ObjectNode) { @@ -48,6 +52,10 @@ public String getRid() { return this.rid; } + public long getDocumentRid() { + return this.documentRid; + } + public void setRid(String rid) { super.setResourceId(rid); } From 25280163dafc043cc9fcef92216caa6113d6f79c Mon Sep 17 00:00:00 2001 From: Maya Painter Date: Mon, 26 Oct 2026 11:35:07 -0700 Subject: [PATCH 2/5] Fix hybrid search integration assertions for ordinal RRF Derive strict fusion expectations from standalone component queries and numeric document RIDs. Allow only zero-score title ties to vary in non-RRF queries, and use cleanup-compliant test database creation. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 29c3522f-9887-43c5-bebb-f42ee9ba4d51 --- .../cosmos/rx/HybridSearchQueryTest.java | 145 +++++++++++++++--- .../test-resource-hygiene-baseline.properties | 1 - 2 files changed, 126 insertions(+), 20 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/HybridSearchQueryTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/HybridSearchQueryTest.java index 6b4ba0e251f1e..848ab1c8cf0ba 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/HybridSearchQueryTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/HybridSearchQueryTest.java @@ -8,6 +8,7 @@ import com.azure.cosmos.CosmosAsyncDatabase; import com.azure.cosmos.CosmosClientBuilder; import com.azure.cosmos.CosmosException; +import com.azure.cosmos.implementation.ResourceId; import com.azure.cosmos.implementation.TestConfigurations; import com.azure.cosmos.implementation.query.HybridSearchBadRequestException; import com.azure.cosmos.models.CosmosContainerProperties; @@ -45,11 +46,13 @@ import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; +import java.util.HashMap; import java.util.List; +import java.util.Map; import java.util.UUID; import java.util.stream.Collectors; -import static com.azure.cosmos.rx.TestSuiteBase.createDatabase; +import static com.azure.cosmos.rx.TestSuiteBase.createTestDatabase; import static com.azure.cosmos.rx.TestSuiteBase.createCollection; import static com.azure.cosmos.rx.TestSuiteBase.safeClose; import static com.azure.cosmos.rx.TestSuiteBase.safeDeleteDatabase; @@ -64,12 +67,19 @@ public class HybridSearchQueryTest { protected static Logger logger = LoggerFactory.getLogger(HybridSearchQueryTest.class); + private static final List JOHN_AND_UNITED_IDS = + Arrays.asList("2", "22", "24", "25", "49", "51", "54", "57", "61", "66", "75", "76", "77", "80", "85"); private final String containerId = UUID.randomUUID().toString(); private final String hpkContainerId = UUID.randomUUID().toString(); private CosmosAsyncClient client; private CosmosAsyncDatabase database; private CosmosAsyncContainer container; private CosmosAsyncContainer hpkContainer; + private final Map documentRids = new HashMap<>(); + private List titleJohnOrder; + private List textJohnOrder; + private List textUnitedOrder; + private List vectorOrder; @BeforeClass(groups = {"query", "split"}, timeOut = SETUP_TIMEOUT* 10) public void before_HybridSearchQueryTest() { @@ -79,7 +89,7 @@ public void before_HybridSearchQueryTest() { .key(TestConfigurations.MASTER_KEY) .buildAsyncClient(); - database = createDatabase(client, UUID.randomUUID().toString()); + database = createTestDatabase(client, "hybrid-search"); PartitionKeyDefinition partitionKeyDef = new PartitionKeyDefinition(); ArrayList paths = new ArrayList(); @@ -129,6 +139,36 @@ public void before_HybridSearchQueryTest() { doc.category = (index % 3 == 0) ? "A" : "B"; hpkContainer.createItem(doc).block(); } + + List storedDocuments = container.queryItems("SELECT c.id, c._rid FROM c", + new CosmosQueryRequestOptions(), Document.class).collectList().block(); + assertThat(storedDocuments).hasSize(100); + for (Document doc : storedDocuments) { + documentRids.put(doc.id, ResourceId.parse(doc.rid).getDocument()); + } + + // Single-component queries bypass client-side RRF. Keep their service score ordering, + // but explicitly order the fixture's zero-score ties by numeric RID for ordinal fusion. + titleJohnOrder = fullTextOrder("title", "John", Collections.singletonList("2")); + textJohnOrder = fullTextOrder("text", "John", Arrays.asList("2", "57", "85")); + textUnitedOrder = fullTextOrder("text", "United", + Arrays.asList("22", "24", "25", "49", "51", "54", "61", "66", "75", "76", "77", "80")); + SqlQuerySpec vectorQuery = new SqlQuerySpec( + "SELECT TOP 100 c.id, VectorDistance(c.vector, @vector) AS score FROM c " + + "ORDER BY VectorDistance(c.vector, @vector)", + Collections.singletonList(new SqlParameter("@vector", getQueryVector()))); + List vectorDocuments = container.queryItems(vectorQuery, + new CosmosQueryRequestOptions(), Document.class).collectList().block(); + assertThat(vectorDocuments).hasSize(100); + // The default cosine VectorDistance projects similarity, not a lower-is-better distance. + assertThat(vectorDocuments).extracting(document -> document.score) + .isSortedAccordingTo(Collections.reverseOrder()); + vectorDocuments.sort((left, right) -> { + int comparison = Double.compare(right.score, left.score); + return comparison != 0 ? comparison + : Long.compareUnsigned(documentRids.get(right.id), documentRids.get(left.id)); + }); + vectorOrder = vectorDocuments.stream().map(Document::getId).collect(Collectors.toList()); } @AfterClass(groups = {"query", "split"}, timeOut = SHUTDOWN_TIMEOUT, alwaysRun = true) @@ -149,7 +189,7 @@ public void hybridQueryTest() { .flatMap(feedResponse -> Flux.fromIterable(feedResponse.getResults())) .collectList().block(); assertThat(resultDocs).hasSize(3); - validateResults(Arrays.asList("2","57","85"), resultDocs); + validateTitleJohnResults(resultDocs, true); query = "SELECT c.id, c.title FROM c WHERE FullTextContains(c.title, @term1) " + "OR FullTextContains(c.text, @term1) ORDER BY RANK FullTextScore(c.title, @term1) OFFSET 1 LIMIT 5"; @@ -160,7 +200,7 @@ public void hybridQueryTest() { .flatMap(feedResponse -> Flux.fromIterable(feedResponse.getResults())) .collectList().block(); assertThat(resultDocs).hasSize(2); - validateResults(Arrays.asList("57","85"), resultDocs); + validateTitleJohnResults(resultDocs, false); query = "SELECT TOP 20 c.id, c.title FROM c WHERE FullTextContains(c.title, @term1) OR " + "FullTextContains(c.text, @term1) OR FullTextContains(c.text, @term2) " + @@ -173,7 +213,9 @@ public void hybridQueryTest() { .flatMap(feedResponse -> Flux.fromIterable(feedResponse.getResults())) .collectList().block(); assertThat(resultDocs).hasSize(15); - validateResults(Arrays.asList("61", "54", "51", "49", "24", "2", "57", "22", "75", "25", "77", "76", "66", "80", "85"), resultDocs); + List expected = expectedRrfOrder(JOHN_AND_UNITED_IDS, new double[] {1, 1}, + titleJohnOrder, textUnitedOrder); + validateResults(expected, resultDocs); query = "SELECT c.id, c.title FROM c WHERE FullTextContains(c.title, @term1) " + "OR FullTextContains(c.text, @term1) OR FullTextContains(c.text, @term2) ORDER BY " + @@ -186,7 +228,7 @@ public void hybridQueryTest() { .flatMap(feedResponse -> Flux.fromIterable(feedResponse.getResults())) .collectList().block(); assertThat(resultDocs).hasSize(10); - validateResults(Arrays.asList("2", "57", "22", "75", "25", "77", "76", "66", "80", "85"), resultDocs); + validateResults(expected.subList(5, 15), resultDocs); List vector = getQueryVector(); query = "SELECT TOP 10 c.id, c.text, c.title FROM c " + @@ -201,7 +243,8 @@ public void hybridQueryTest() { .flatMap(feedResponse -> Flux.fromIterable(feedResponse.getResults())) .collectList().block(); assertThat(resultDocs).hasSize(10); - validateResults(Arrays.asList("55", "61", "57", "24", "2", "54", "63", "9", "62", "75"), resultDocs); + validateResults(expectedRrfOrder(new ArrayList<>(documentRids.keySet()), new double[] {1, 1, 1}, + textJohnOrder, textUnitedOrder, vectorOrder).subList(0, 10), resultDocs); } @Test(groups = {"query", "split"}, timeOut = TIMEOUT) @@ -221,7 +264,7 @@ public void hybridQueryWeightedRRFTest(){ assertThat(results).isNotNull(); validateResults( - Arrays.asList("61", "54", "51", "49", "24", "2", "57", "22", "75", "25", "77", "76", "66", "80", "85"), + expectedRrfOrder(JOHN_AND_UNITED_IDS, new double[] {1, 1}, titleJohnOrder, textUnitedOrder), results ); @@ -240,7 +283,7 @@ public void hybridQueryWeightedRRFTest(){ assertThat(results).hasSize(15); assertThat(results).isNotNull(); validateResults( - Arrays.asList("61", "54", "51", "49", "24", "2", "57", "22", "75", "25", "77", "76", "66", "80", "85"), + expectedRrfOrder(JOHN_AND_UNITED_IDS, new double[] {10, 10}, titleJohnOrder, textUnitedOrder), results ); @@ -258,7 +301,8 @@ public void hybridQueryWeightedRRFTest(){ assertThat(results).hasSize(10); assertThat(results).isNotNull(); validateResults( - Arrays.asList("61", "54", "51", "49", "24", "2", "57", "22", "75", "25"), + expectedRrfOrder(JOHN_AND_UNITED_IDS, new double[] {0.1, 0.1}, + titleJohnOrder, textUnitedOrder).subList(0, 10), results ); @@ -276,7 +320,8 @@ public void hybridQueryWeightedRRFTest(){ assertThat(results).hasSize(10); assertThat(results).isNotNull(); validateResults( - Arrays.asList("57", "22", "25", "54", "66", "24", "2", "85", "61", "76"), + expectedRrfOrder(JOHN_AND_UNITED_IDS, new double[] {-1, -1}, + titleJohnOrder, textUnitedOrder).subList(0, 10), results ); @@ -296,7 +341,8 @@ public void hybridQueryWeightedRRFTest(){ assertThat(results).hasSize(10); assertThat(results).isNotNull(); validateResults( - Arrays.asList("75", "24", "49", "55", "61", "21", "9", "26", "37", "57"), + expectedRrfOrder(new ArrayList<>(documentRids.keySet()), new double[] {1, 1}, + textUnitedOrder, vectorOrder).subList(0, 10), results ); } @@ -328,9 +374,8 @@ public void hybridQueryWithLocalStatisticsTest() { assertThat(explicitGlobalResultDocs).isNotNull(); assertThat(explicitGlobalResultDocs).hasSize(3); - List defaultIds = globalResultDocs.stream().map(Document::getId).collect(Collectors.toList()); - List explicitGlobalIds = explicitGlobalResultDocs.stream().map(Document::getId).collect(Collectors.toList()); - assertThat(explicitGlobalIds).isEqualTo(defaultIds); + validateTitleJohnResults(globalResultDocs, true); + validateTitleJohnResults(explicitGlobalResultDocs, true); // Test 3: LOCAL scope with pk="2" — only id=57 and id=85 have 'John' in pk="2" // Stats are computed only over pk="2" partition @@ -346,7 +391,7 @@ public void hybridQueryWithLocalStatisticsTest() { for (Document doc : localPk2ResultDocs) { assertThat(doc.getPk()).isEqualTo("2"); } - validateResults(Arrays.asList("57", "85"), localPk2ResultDocs); + validateTitleJohnResults(localPk2ResultDocs, false); // Test 4: LOCAL scope with pk="1" — only id=2 has 'John' in pk="1" // Stats are computed only over pk="1" partition @@ -500,7 +545,7 @@ public void hybridQueryWithLocalScopeAndFeedRangeTest() { for (Document doc : feedRangeResults) { assertThat(doc.getPk()).isEqualTo("2"); } - validateResults(Arrays.asList("57", "85"), feedRangeResults); + validateTitleJohnResults(feedRangeResults, false); // LOCAL scope with FeedRange for pk="1" — only id=2 matches CosmosQueryRequestOptions localFeedRangePk1 = new CosmosQueryRequestOptions(); @@ -648,11 +693,71 @@ private static List loadProductsFromJson() { } } - private void validateResults(List actualIds, List results) { + private List fullTextOrder(String field, String term, List matchingIds) { + SqlQuerySpec query = new SqlQuerySpec( + "SELECT TOP 100 c.id FROM c ORDER BY RANK FullTextScore(c." + field + ", @term)", + Collections.singletonList(new SqlParameter("@term", term))); + List serviceOrder = container.queryItems(query, new CosmosQueryRequestOptions(), Document.class) + .map(Document::getId).collectList().block(); + assertThat(serviceOrder).containsExactlyInAnyOrderElementsOf(documentRids.keySet()); + assertThat(serviceOrder.subList(0, matchingIds.size())).containsExactlyInAnyOrderElementsOf(matchingIds); + + List order = new ArrayList<>(serviceOrder.subList(0, matchingIds.size())); + List zeroScores = serviceOrder.stream().filter(id -> !matchingIds.contains(id)) + .sorted((left, right) -> Long.compareUnsigned(documentRids.get(right), documentRids.get(left))) + .collect(Collectors.toList()); + order.addAll(zeroScores); + return order; + } + + @SafeVarargs + private final List expectedRrfOrder(List candidateIds, double[] weights, + List... componentOrders) { + Map scores = new HashMap<>(); + for (String id : candidateIds) { + scores.put(id, 0.0); + } + + // Rank only the query's candidates. Negative weights reverse the component direction + // (including RID ties); the plan normalizes their magnitude before applying RRF. + for (int component = 0; component < componentOrders.length; component++) { + List order = componentOrders[component].stream().filter(candidateIds::contains) + .collect(Collectors.toList()); + assertThat(order).containsExactlyInAnyOrderElementsOf(candidateIds); + if (weights[component] < 0) { + Collections.reverse(order); + } + for (int position = 0; position < order.size(); position++) { + String id = order.get(position); + scores.put(id, scores.get(id) + Math.abs(weights[component]) / (60 + position + 1)); + } + } + + List expected = new ArrayList<>(candidateIds); + expected.sort((left, right) -> { + int comparison = Double.compare(scores.get(right), scores.get(left)); + return comparison != 0 ? comparison + : Long.compareUnsigned(documentRids.get(right), documentRids.get(left)); + }); + return expected; + } + + private void validateTitleJohnResults(List results, boolean includesTitleMatch) { + List ids = results.stream().map(Document::getId).collect(Collectors.toList()); + // Only id 2 has "John" in its title. IDs 57/85 match text, but both title scores are zero. + // A single FullTextScore query bypasses RRF, so it does not promise the RRF RID tie order. + if (includesTitleMatch) { + assertThat(ids.get(0)).isEqualTo("2"); + ids = ids.subList(1, ids.size()); + } + assertThat(ids).containsExactlyInAnyOrder("57", "85"); + } + + private void validateResults(List expectedIds, List results) { assertThat(results).isNotNull(); List resultsIds = results.stream().map(Document::getId).collect(Collectors.toList()); - assertThat(resultsIds).isEqualTo(actualIds); + assertThat(resultsIds).isEqualTo(expectedIds); } public static List getQueryVector() { @@ -666,6 +771,8 @@ public static List getQueryVector() { static class Document { @JsonProperty("id") String id; + @JsonProperty("_rid") + String rid; @JsonProperty("pk") String pk; @JsonProperty("category") diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/resources/test-resource-hygiene-baseline.properties b/sdk/cosmos/azure-cosmos-tests/src/test/resources/test-resource-hygiene-baseline.properties index 22bdae8da6383..9f80bad59795c 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/resources/test-resource-hygiene-baseline.properties +++ b/sdk/cosmos/azure-cosmos-tests/src/test/resources/test-resource-hygiene-baseline.properties @@ -40,7 +40,6 @@ com/azure/cosmos/rx/FullTextIndexTest.java=1 com/azure/cosmos/rx/GatewayReadConsistencyStrategyE2ETest.java=1 com/azure/cosmos/rx/GatewayReadConsistencyStrategySpyWireTest.java=1 com/azure/cosmos/rx/GlobalSecondaryIndexContainerCrudTest.java=1 -com/azure/cosmos/rx/HybridSearchQueryTest.java=1 com/azure/cosmos/rx/MultiMasterConflictResolutionTest.java=1 com/azure/cosmos/rx/NonStreamingOrderByQueryVectorSearchTest.java=1 com/azure/cosmos/rx/OfferQueryTest.java=1 From 2cdf0671bf700c8f5c4642e1cc3d40a5b8265c34 Mon Sep 17 00:00:00 2001 From: Maya Painter Date: Mon, 26 Oct 2026 12:57:43 -0700 Subject: [PATCH 3/5] Apply hybrid search test review formatting Keep the full-text reference setup and numeric RID tie-break expression on single lines as requested. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 29c3522f-9887-43c5-bebb-f42ee9ba4d51 --- .../java/com/azure/cosmos/rx/HybridSearchQueryTest.java | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/HybridSearchQueryTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/HybridSearchQueryTest.java index 848ab1c8cf0ba..c857412d02938 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/HybridSearchQueryTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/HybridSearchQueryTest.java @@ -151,8 +151,7 @@ public void before_HybridSearchQueryTest() { // but explicitly order the fixture's zero-score ties by numeric RID for ordinal fusion. titleJohnOrder = fullTextOrder("title", "John", Collections.singletonList("2")); textJohnOrder = fullTextOrder("text", "John", Arrays.asList("2", "57", "85")); - textUnitedOrder = fullTextOrder("text", "United", - Arrays.asList("22", "24", "25", "49", "51", "54", "61", "66", "75", "76", "77", "80")); + textUnitedOrder = fullTextOrder("text", "United", Arrays.asList("22", "24", "25", "49", "51", "54", "61", "66", "75", "76", "77", "80")); SqlQuerySpec vectorQuery = new SqlQuerySpec( "SELECT TOP 100 c.id, VectorDistance(c.vector, @vector) AS score FROM c " + "ORDER BY VectorDistance(c.vector, @vector)", @@ -165,8 +164,7 @@ public void before_HybridSearchQueryTest() { .isSortedAccordingTo(Collections.reverseOrder()); vectorDocuments.sort((left, right) -> { int comparison = Double.compare(right.score, left.score); - return comparison != 0 ? comparison - : Long.compareUnsigned(documentRids.get(right.id), documentRids.get(left.id)); + return comparison != 0 ? comparison : Long.compareUnsigned(documentRids.get(right.id), documentRids.get(left.id)); }); vectorOrder = vectorDocuments.stream().map(Document::getId).collect(Collectors.toList()); } From 093a7cd5507a39edf930608efa53948a1dbf4e47 Mon Sep 17 00:00:00 2001 From: Maya Painter Date: Mon, 26 Oct 2026 13:39:18 -0700 Subject: [PATCH 4/5] Test hybrid ranking through emulator queries Restore private ranking helpers and replace implementation-coupled unit tests with public-query integration coverage for ordinal and competition ranks, numeric RID ties, weights, and paging. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 29c3522f-9887-43c5-bebb-f42ee9ba4d51 --- .../query/HybridSearchRankingTest.java | 295 ------------------ .../cosmos/rx/HybridSearchRankingTest.java | 254 +++++++++++++++ ...idSearchDocumentQueryExecutionContext.java | 12 +- 3 files changed, 260 insertions(+), 301 deletions(-) delete mode 100644 sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/query/HybridSearchRankingTest.java create mode 100644 sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/HybridSearchRankingTest.java diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/query/HybridSearchRankingTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/query/HybridSearchRankingTest.java deleted file mode 100644 index 7e09592338203..0000000000000 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/query/HybridSearchRankingTest.java +++ /dev/null @@ -1,295 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -package com.azure.cosmos.implementation.query; - -import com.azure.cosmos.implementation.Configs; -import com.azure.cosmos.implementation.Document; -import com.azure.cosmos.implementation.ResourceType; -import com.azure.cosmos.implementation.query.HybridSearchDocumentQueryExecutionContext.ComponentWeight; -import com.azure.cosmos.implementation.query.HybridSearchDocumentQueryExecutionContext.ScoreTuple; -import com.azure.cosmos.implementation.query.hybridsearch.HybridSearchQueryInfo; -import com.azure.cosmos.implementation.query.hybridsearch.HybridSearchQueryResult; -import com.azure.cosmos.models.CosmosQueryRequestOptions; -import com.azure.cosmos.models.FeedResponse; -import com.azure.cosmos.models.SqlQuerySpec; -import org.testng.annotations.DataProvider; -import org.testng.annotations.Test; -import reactor.core.publisher.Flux; -import reactor.core.publisher.Mono; - -import java.lang.reflect.Field; -import java.nio.ByteBuffer; -import java.nio.ByteOrder; -import java.util.ArrayList; -import java.util.Arrays; -import java.util.Base64; -import java.util.Collections; -import java.util.List; -import java.util.UUID; -import java.util.concurrent.atomic.AtomicBoolean; -import java.util.concurrent.atomic.AtomicInteger; - -import static com.azure.cosmos.implementation.query.HybridSearchDocumentQueryExecutionContext.coalesceAndSortResults; -import static com.azure.cosmos.implementation.query.HybridSearchDocumentQueryExecutionContext.computeRRFScores; -import static com.azure.cosmos.implementation.query.HybridSearchDocumentQueryExecutionContext.computeRanks; -import static com.azure.cosmos.implementation.query.HybridSearchDocumentQueryExecutionContext.retrieveComponentScores; -import static com.azure.cosmos.implementation.query.HybridSearchDocumentQueryExecutionContext.retrieveComponentWeights; -import static org.assertj.core.api.Assertions.assertThat; -import static org.assertj.core.api.Assertions.within; -import static org.mockito.Mockito.doReturn; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.spy; - -public class HybridSearchRankingTest { - private static final String RANKING_PROPERTY = "COSMOS.HYBRID_SEARCH_USE_COMPETITION_RANKING"; - - @DataProvider(name = "rankingModes") - public Object[][] rankingModes() { - return new Object[][] { - {false, Arrays.asList(1, 2, 3, 4, 5)}, - {true, Arrays.asList(1, 2, 2, 2, 5)} - }; - } - - @Test(groups = "unit", dataProvider = "rankingModes") - public void tiedScoresUseOrdinalOrCompetitionRanks(boolean competition, List expected) { - List> rows = Arrays.asList( - row(5, 10), row(4, 5), row(3, 5), row(2, 5), row(1, 0)); - List weights = weights(null, SortOrder.Descending); - - List> ranks = ranks(rows, weights, competition); - assertThat(ranks).containsExactly(expected); - - List> fused = computeRRFScores( - Mono.just(ranks), Mono.just(new ArrayList<>(rows)), weights).collectList().block(); - assertThat(fused).containsExactlyElementsOf(rows); - for (int i = 0; i < rows.size(); i++) { - assertThat(rows.get(i).getScore()).isCloseTo(1.0 / (60 + expected.get(i)), within(1e-15)); - } - } - - @DataProvider(name = "directions") - public Object[][] directions() { - return new Object[][] { - {SortOrder.Ascending, Arrays.asList(1, 2)}, - {SortOrder.Descending, Arrays.asList(2, 1)} - }; - } - - @Test(groups = "unit", dataProvider = "directions") - public void tiesFollowPlanDirectionRatherThanEncodedRidOrWeight( - SortOrder direction, List expected) { - HybridSearchQueryResult low = row(1, 5); - HybridSearchQueryResult high = row(256, 5); - assertThat(low.getRid().compareTo(high.getRid())).isPositive(); - - // A negative original weight can arrive normalized to positive with an ascending plan. - List weights = weights(Collections.singletonList(2.0), direction); - assertThat(ranks(Arrays.asList(low, high), weights, false)).containsExactly(expected); - assertThat(ranks(Arrays.asList(high, low), weights, false)) - .containsExactly(Arrays.asList(expected.get(1), expected.get(0))); - assertThat(ranks(Arrays.asList(low, high), weights, true)) - .containsExactly(Arrays.asList(1, 1)); - } - - @Test(groups = "unit") - public void mixedDirectionsPreserveWeightsAndFormula() { - List> rows = Arrays.asList(row(1, 10, 5), row(256, 5, 10)); - List weights = weights(Arrays.asList(2.0, 3.0), - SortOrder.Descending, SortOrder.Ascending); - List> ranks = ranks(rows, weights, false); - - assertThat(ranks).containsExactly(Arrays.asList(1, 2), Arrays.asList(1, 2)); - computeRRFScores(Mono.just(ranks), Mono.just(rows), weights).collectList().block(); - assertThat(rows.get(0).getScore()).isCloseTo(5.0 / 61, within(1e-15)); - assertThat(rows.get(1).getScore()).isCloseTo(5.0 / 62, within(1e-15)); - } - - @Test(groups = "unit") - public void finalScoreTiesUseDescendingNumericDocumentRid() { - List> rows = Arrays.asList(row(1, 10, 5), row(256, 5, 10)); - List weights = weights(null, SortOrder.Descending, SortOrder.Descending); - List> fused = computeRRFScores( - Mono.just(ranks(rows, weights, false)), Mono.just(rows), weights).collectList().block(); - - assertThat(fused).extracting(HybridSearchQueryResult::getDocumentRid).containsExactly(256L, 1L); - assertThat(fused.get(0).getScore()).isEqualTo(fused.get(1).getScore()); - } - - @Test(groups = "unit") - public void ridDecodingUsesLittleEndianAndCosmosBase64() { - long documentRid = 0x000102030405ffffL; - HybridSearchQueryResult result = row(documentRid, 5); - - assertThat(result.getRid()).contains("-"); - assertThat(result.getDocumentRid()).isEqualTo(documentRid); - assertThat(row(128, 5).getDocumentRid()).isEqualTo(128); - } - - @Test(groups = "unit") - public void numericRidIsCachedBeforeRankingAndFusion() { - HybridSearchQueryResult low = row(1, 5); - HybridSearchQueryResult high = row(256, 5); - low.rid = "not-a-rid"; - high.rid = "not-a-rid-either"; - List> rows = Arrays.asList(low, high); - List weights = weights(null, SortOrder.Descending); - - assertThat(low.getDocumentRid()).isEqualTo(1); - assertThat(high.getDocumentRid()).isEqualTo(256); - assertThat(ranks(rows, weights, false)).containsExactly(Arrays.asList(2, 1)); - List> fused = computeRRFScores( - Mono.just(ranks(rows, weights, true)), Mono.just(rows), weights).collectList().block(); - assertThat(fused).containsExactly(high, low); - } - - @Test(groups = "unit") - public void numericRidComparisonsAreUnsigned() { - HybridSearchQueryResult high = spy(row(256, 5)); - doReturn(-1L).when(high).getDocumentRid(); - HybridSearchQueryResult low = row(1, 5); - List> rows = Arrays.asList(low, high); - List weights = weights(null, SortOrder.Descending); - - assertThat(ranks(rows, weights, false)).containsExactly(Arrays.asList(2, 1)); - List> fused = computeRRFScores( - Mono.just(ranks(rows, weights, true)), Mono.just(rows), weights).collectList().block(); - assertThat(fused).containsExactly(high, low); - } - - @Test(groups = "unit") - public void deduplicationKeepsEncodedIdentityAndFirstPayload() { - String rid = rid(1, 1); - Document first = document(rid, "first", 10); - Document duplicate = document(rid, "duplicate", 100); - Document otherCollection = document(rid(2, 1), "other-collection", 5); - Document high = document(rid(1, 256), "high", 3); - - List> rows = coalesceAndSortResults( - Flux.just(first, duplicate, otherCollection, high)).block(); - - assertThat(rows).hasSize(3); - assertThat(rows).extracting(HybridSearchQueryResult::getDocumentRid).containsExactly(256L, 1L, 1L); - assertThat(rows).extracting(result -> result.getPayload().getId()) - .containsExactly("high", "first", "other-collection"); - assertThat(rows.get(1).getComponentScores()).containsExactly(10.0); - } - - @Test(groups = "unit", dataProvider = "rankingModes") - public void emptyScoresUsePlanDimensionsAndExistingUndefinedValue( - boolean competition, List ignored) { - List> rows = Arrays.asList(row(256), row(1, 10, 5)); - List weights = weights(null, SortOrder.Descending, SortOrder.Ascending); - List> scores = retrieveComponentScores(Mono.just(rows), weights).block(); - - assertThat(scores).hasSize(2); - assertThat(scores.get(0)).extracting(ScoreTuple::getScore).containsExactly(10.0, -999999.0); - assertThat(scores.get(1)).extracting(ScoreTuple::getScore).containsExactly(-999999.0, 5.0); - assertThat(ranks(rows, weights, competition)).containsExactly(Arrays.asList(2, 1), Arrays.asList(1, 2)); - - Collections.reverse(rows); - assertThat(ranks(rows, weights, competition)).containsExactly(Arrays.asList(1, 2), Arrays.asList(2, 1)); - } - - @Test(groups = "unit") - public void emptyInputAndAllEmptyScoresAreSupported() { - List weights = weights(null, SortOrder.Descending, SortOrder.Descending); - Mono>> empty = coalesceAndSortResults(Flux.empty()).cache(); - Mono>> emptyRanks = computeRanks(retrieveComponentScores(empty, weights), false); - assertThat(computeRRFScores(emptyRanks, empty, weights).collectList().block()).isEmpty(); - - List> rows = Arrays.asList(row(256), row(1)); - assertThat(ranks(rows, weights, false)).containsExactly(Arrays.asList(1, 2), Arrays.asList(1, 2)); - assertThat(ranks(rows, weights, true)).containsExactly(Arrays.asList(1, 1), Arrays.asList(1, 1)); - assertThat(ranks(Collections.singletonList(row(1)), weights, false)) - .containsExactly(Collections.singletonList(1), Collections.singletonList(1)); - } - - @Test(groups = "unit") - public void rankingReusesCollectedRowsAndPreservesPayloadPaging() throws ReflectiveOperationException { - AtomicInteger subscriptions = new AtomicInteger(); - Mono>> rows = coalesceAndSortResults(Flux.defer(() -> { - subscriptions.incrementAndGet(); - return Flux.just(document(rid(1, 1), "low", 5), document(rid(1, 256), "high", 5), - document(rid(1, 128), "middle", 5)); - })).cache(); - List weights = weights(null, SortOrder.Descending); - Flux> fused = computeRRFScores( - computeRanks(retrieveComponentScores(rows, weights), false), rows, weights); - HybridSearchDocumentQueryExecutionContext context = context(); - Field observable = HybridSearchDocumentQueryExecutionContext.class.getDeclaredField("hybridObservable"); - observable.setAccessible(true); - observable.set(context, fused); - - List> pages = context.drainAsync(2).collectList().block(); - - assertThat(subscriptions.get()).isEqualTo(1); - assertThat(pages).hasSize(2); - assertThat(pages.get(0).getResults()).extracting(Document::getId).containsExactly("high", "middle"); - assertThat(pages.get(1).getResults()).extracting(Document::getId).containsExactly("low"); - } - - @Test(groups = "unit") - public void configurationIsResolvedOncePerQuery() throws ReflectiveOperationException { - String previous = System.getProperty(RANKING_PROPERTY); - try { - System.clearProperty(RANKING_PROPERTY); - assertThat(Configs.isHybridSearchCompetitionRankingEnabled()).isEqualTo( - Boolean.parseBoolean(System.getenv("AZURE_COSMOS_HYBRID_SEARCH_USE_COMPETITION_RANKING"))); - - System.setProperty(RANKING_PROPERTY, "false"); - HybridSearchDocumentQueryExecutionContext ordinal = context(); - System.setProperty(RANKING_PROPERTY, "true"); - HybridSearchDocumentQueryExecutionContext competition = context(); - Field rankingFlag = HybridSearchDocumentQueryExecutionContext.class.getDeclaredField("useCompetitionRanking"); - rankingFlag.setAccessible(true); - assertThat(rankingFlag.getBoolean(ordinal)).isFalse(); - assertThat(rankingFlag.getBoolean(competition)).isTrue(); - } finally { - if (previous == null) { - System.clearProperty(RANKING_PROPERTY); - } else { - System.setProperty(RANKING_PROPERTY, previous); - } - } - } - - private static HybridSearchDocumentQueryExecutionContext context() { - return new HybridSearchDocumentQueryExecutionContext(null, mock(IDocumentQueryClient.class), - ResourceType.Document, new SqlQuerySpec("SELECT * FROM c"), new CosmosQueryRequestOptions(), - "dbs/db/colls/coll", null, UUID.randomUUID(), new AtomicBoolean(), new HybridSearchQueryInfo()); - } - - private static List> ranks(List> rows, - List weights, boolean competition) { - return computeRanks(retrieveComponentScores(Mono.just(rows), weights), competition).block(); - } - - private static List weights(List values, SortOrder... directions) { - List components = new ArrayList<>(); - for (SortOrder direction : directions) { - QueryInfo info = new QueryInfo(); - info.setOrderBy(Collections.singletonList(direction)); - components.add(info); - } - return retrieveComponentWeights(values, components); - } - - private static HybridSearchQueryResult row(long documentRid, double... scores) { - return new HybridSearchQueryResult<>(document(rid(1, documentRid), "doc-" + documentRid, scores).toJson()); - } - - private static Document document(String rid, String id, double... scores) { - return new Document("{\"_rid\":\"" + rid + "\",\"payload\":{\"componentScores\":" - + Arrays.toString(scores) + ",\"payload\":{\"id\":\"" + id + "\"}}}"); - } - - private static String rid(int database, long document) { - byte[] bytes = new byte[16]; - ByteBuffer.wrap(bytes).putInt(database); - bytes[4] = (byte) 0x80; - ByteBuffer.wrap(bytes, 8, 8).order(ByteOrder.LITTLE_ENDIAN).putLong(document); - return Base64.getEncoder().encodeToString(bytes).replace('/', '-'); - } -} diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/HybridSearchRankingTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/HybridSearchRankingTest.java new file mode 100644 index 0000000000000..0e1bb6084600e --- /dev/null +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/HybridSearchRankingTest.java @@ -0,0 +1,254 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.rx; + +import com.azure.cosmos.CosmosAsyncClient; +import com.azure.cosmos.CosmosAsyncContainer; +import com.azure.cosmos.CosmosAsyncDatabase; +import com.azure.cosmos.CosmosClientBuilder; +import com.azure.cosmos.implementation.TestConfigurations; +import com.azure.cosmos.models.CosmosContainerProperties; +import com.azure.cosmos.models.CosmosContainerRequestOptions; +import com.azure.cosmos.models.CosmosFullTextIndex; +import com.azure.cosmos.models.CosmosFullTextPath; +import com.azure.cosmos.models.CosmosFullTextPolicy; +import com.azure.cosmos.models.CosmosQueryRequestOptions; +import com.azure.cosmos.models.FeedResponse; +import com.azure.cosmos.models.IndexingPolicy; +import com.fasterxml.jackson.databind.node.JsonNodeFactory; +import com.fasterxml.jackson.databind.node.ObjectNode; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.DataProvider; +import org.testng.annotations.Test; + +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Base64; +import java.util.List; +import java.util.stream.Collectors; + +import static com.azure.cosmos.rx.TestSuiteBase.createCollection; +import static com.azure.cosmos.rx.TestSuiteBase.createTestDatabase; +import static com.azure.cosmos.rx.TestSuiteBase.safeClose; +import static com.azure.cosmos.rx.TestSuiteBase.safeDeleteDatabase; +import static org.assertj.core.api.Assertions.assertThat; + +// Runs in the query integration suite against an emulator/account with full-text search support. +@Test(groups = "query", singleThreaded = true) +public class HybridSearchRankingTest { + private static final String RANKING_PROPERTY = "COSMOS.HYBRID_SEARCH_USE_COMPETITION_RANKING"; + private static final int DOCUMENT_COUNT = 260; + private CosmosAsyncClient client; + private CosmosAsyncDatabase database; + private CosmosAsyncContainer container; + private CosmosAsyncContainer tierContainer; + private List documentsByRid; + + @BeforeClass(groups = "query", timeOut = 120000) + public void setUp() { + client = new CosmosClientBuilder().endpoint(TestConfigurations.HOST) + .key(TestConfigurations.MASTER_KEY).buildAsyncClient(); + database = createTestDatabase(client, "hybrid-ranking"); + + CosmosFullTextIndex titleIndex = new CosmosFullTextIndex(); + titleIndex.setPath("/title"); + CosmosFullTextIndex textIndex = new CosmosFullTextIndex(); + textIndex.setPath("/text"); + IndexingPolicy indexing = new IndexingPolicy(); + indexing.setCosmosFullTextIndexes(Arrays.asList(titleIndex, textIndex)); + CosmosFullTextPath titlePath = new CosmosFullTextPath(); + titlePath.setPath("/title"); + titlePath.setLanguage("en-US"); + CosmosFullTextPath textPath = new CosmosFullTextPath(); + textPath.setPath("/text"); + textPath.setLanguage("en-US"); + CosmosFullTextPolicy fullText = new CosmosFullTextPolicy(); + fullText.setDefaultLanguage("en-US"); + fullText.setPaths(Arrays.asList(titlePath, textPath)); + CosmosContainerProperties properties = new CosmosContainerProperties("ties", "/pk"); + properties.setIndexingPolicy(indexing); + properties.setFullTextPolicy(fullText); + container = createCollection(database, properties, new CosmosContainerRequestOptions(), 10000); + CosmosContainerProperties tierProperties = new CosmosContainerProperties("tiers", "/pk"); + tierProperties.setIndexingPolicy(indexing); + tierProperties.setFullTextPolicy(fullText); + tierContainer = createCollection(database, tierProperties, new CosmosContainerRequestOptions(), 400); + String[] tierTexts = {"beta filler filler filler filler", "beta beta filler filler filler", + "beta beta beta filler filler", "beta beta beta beta filler", "beta beta beta beta beta"}; + for (int i = 0; i < 5; i++) { + ObjectNode document = JsonNodeFactory.instance.objectNode(); + document.put("id", "tier-" + i); + document.put("pk", "tiers"); + document.put("title", i == 0 ? "alpha alpha" : i == 4 ? "filler filler" : "alpha filler"); + document.put("text", tierTexts[i]); + tierContainer.createItem(document).block(); + } + + // Identical fields guarantee tied component scores without reproducing the service's BM25 implementation. + // More than 256 inserts also exercise RID byte boundaries rather than assuming encoded RID order. + for (int i = 0; i < DOCUMENT_COUNT; i++) { + ObjectNode document = JsonNodeFactory.instance.objectNode(); + document.put("id", "doc-" + i); + document.put("pk", Integer.toString(i % 2)); + document.put("title", "alpha"); + document.put("text", "beta"); + container.createItem(document).block(); + } + documentsByRid = container.queryItems("SELECT * FROM c", new CosmosQueryRequestOptions(), ObjectNode.class) + .collectList().block(); + assertThat(documentsByRid).hasSize(DOCUMENT_COUNT); + documentsByRid.sort((left, right) -> Long.compareUnsigned(documentRid(right), documentRid(left))); + List encodedOrder = new ArrayList<>(documentsByRid); + encodedOrder.sort((left, right) -> right.get("_rid").asText().compareTo(left.get("_rid").asText())); + assertThat(ids(encodedOrder)).as("Fixture must distinguish numeric from encoded RID ordering") + .isNotEqualTo(ids(documentsByRid)); + } + + @AfterClass(groups = "query", alwaysRun = true, timeOut = 30000) + public void tearDown() { + safeDeleteDatabase(database); + safeClose(client); + } + + @DataProvider(name = "rankingModes") + public Object[][] rankingModes() { + return new Object[][] {{false}, {true}}; + } + + @Test(groups = "query", dataProvider = "rankingModes", timeOut = 60000) + public void tiedComponentsRespectRankingModeAndDirection(boolean competition) { + String previous = System.getProperty(RANKING_PROPERTY); + try { + System.setProperty(RANKING_PROPERTY, Boolean.toString(competition)); + for (double[] weights : new double[][] {{1, 1}, {-1, -1}, {1, -1}, {2, -1}}) { + List expected = expectedOrder(weights, competition); + String orderBy = orderBy(weights); + List> pages = container.queryItems( + "SELECT TOP 260 * FROM c " + orderBy, new CosmosQueryRequestOptions(), ObjectNode.class) + .byPage(17).collectList().block(); + List actual = pages.stream().flatMap(page -> page.getResults().stream()) + .collect(Collectors.toList()); + assertThat(ids(actual)).as("weights=%s, competition=%s", Arrays.toString(weights), competition) + .containsExactlyElementsOf(expected); + assertThat(pages).hasSizeGreaterThan(1); + assertThat(actual).allSatisfy(document -> { + assertThat(document.get("title").asText()).isEqualTo("alpha"); + assertThat(document.get("text").asText()).isEqualTo("beta"); + }); + + // Keep the full fixture in the service's bounded component candidate lists. + List window = container.queryItems( + "SELECT * FROM c " + orderBy + " OFFSET 3 LIMIT 257", + new CosmosQueryRequestOptions(), ObjectNode.class).collectList().block(); + assertThat(ids(window)).containsExactlyElementsOf(expected.subList(3, DOCUMENT_COUNT)); + } + } finally { + restoreRankingProperty(previous); + } + } + + @Test(groups = "query", timeOut = 30000) + public void competitionRanksSkipTiedPositions() { + List documents = tierContainer.queryItems("SELECT * FROM c", + new CosmosQueryRequestOptions(), ObjectNode.class).collectList().block(); + assertThat(documents).hasSize(5); + documents.sort((left, right) -> Long.compareUnsigned(documentRid(right), documentRid(left))); + List titleOrder = new ArrayList<>(documents); + // Equal document lengths: two occurrences outrank one, and no occurrence scores zero. + titleOrder.sort((left, right) -> Integer.compare(titleTier(left), titleTier(right))); + String previous = System.getProperty(RANKING_PROPERTY); + try { + for (boolean competition : new boolean[] {false, true}) { + System.setProperty(RANKING_PROPERTY, Boolean.toString(competition)); + List expected = new ArrayList<>(documents); + int[] titleRanks = competition ? new int[] {1, 2, 2, 2, 5} : new int[] {1, 2, 3, 4, 5}; + expected.sort((left, right) -> { + double leftScore = 1.0 / (60 + titleRanks[titleOrder.indexOf(left)]) + + 1.0 / (60 + textRank(left)); + double rightScore = 1.0 / (60 + titleRanks[titleOrder.indexOf(right)]) + + 1.0 / (60 + textRank(right)); + int comparison = Double.compare(rightScore, leftScore); + return comparison != 0 ? comparison : Long.compareUnsigned(documentRid(right), documentRid(left)); + }); + List results = tierContainer.queryItems("SELECT TOP 5 * FROM c " + + orderBy(new double[] {1, 1}), new CosmosQueryRequestOptions(), ObjectNode.class) + .collectList().block(); + assertThat(ids(results)).as("competition=%s", competition).containsExactlyElementsOf(ids(expected)); + List topTwo = tierContainer.queryItems("SELECT TOP 2 * FROM c " + + orderBy(new double[] {1, 1}), new CosmosQueryRequestOptions(), ObjectNode.class) + .collectList().block(); + assertThat(ids(topTwo)).containsExactlyElementsOf(ids(expected.subList(0, 2))); + } + } finally { + restoreRankingProperty(previous); + } + } + + private static int titleTier(ObjectNode document) { + String title = document.get("title").asText(); + return "alpha alpha".equals(title) ? 0 : "alpha filler".equals(title) ? 1 : 2; + } + + private static int textRank(ObjectNode document) { + // Texts have equal lengths and one through five occurrences of beta, so all scores are distinct. + return 5 - Integer.parseInt(document.get("id").asText().substring("tier-".length())); + } + + @Test(groups = "query", timeOut = 30000) + public void emptyHybridQueryReturnsNoDocuments() { + List results = container.queryItems( + "SELECT TOP 10 * FROM c WHERE c.id = 'missing' " + orderBy(new double[] {1, 1}), + new CosmosQueryRequestOptions(), ObjectNode.class).collectList().block(); + assertThat(results).isEmpty(); + } + + private List expectedOrder(double[] weights, boolean competition) { + List positions = new ArrayList<>(); + for (int i = 0; i < documentsByRid.size(); i++) { + positions.add(i); + } + // In competition mode every tied score has rank 1. Ordinal rank follows the plan direction; + // negative weights reverse that direction before their absolute magnitude is used in fusion. + positions.sort((left, right) -> { + int comparison = Double.compare(score(right, weights, competition), score(left, weights, competition)); + return comparison != 0 ? comparison : Integer.compare(left, right); + }); + return positions.stream().map(position -> documentsByRid.get(position).get("id").asText()) + .collect(Collectors.toList()); + } + + private static double score(int descendingPosition, double[] weights, boolean competition) { + double score = 0; + for (double weight : weights) { + int rank = competition ? 1 : weight < 0 ? DOCUMENT_COUNT - descendingPosition : descendingPosition + 1; + score += Math.abs(weight) / (60 + rank); + } + return score; + } + + private static String orderBy(double[] weights) { + return "ORDER BY RANK RRF(FullTextScore(c.title, 'alpha'), FullTextScore(c.text, 'beta'), " + + Arrays.toString(weights) + ")"; + } + + private static long documentRid(ObjectNode document) { + byte[] bytes = Base64.getDecoder().decode(document.get("_rid").asText().replace('-', '/')); + return ByteBuffer.wrap(bytes, 8, 8).order(ByteOrder.LITTLE_ENDIAN).getLong(); + } + + private static List ids(List documents) { + return documents.stream().map(document -> document.get("id").asText()).collect(Collectors.toList()); + } + + private static void restoreRankingProperty(String previous) { + if (previous == null) { + System.clearProperty(RANKING_PROPERTY); + } else { + System.setProperty(RANKING_PROPERTY, previous); + } + } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/HybridSearchDocumentQueryExecutionContext.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/HybridSearchDocumentQueryExecutionContext.java index 3431df1f9a542..c7d6661917ed5 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/HybridSearchDocumentQueryExecutionContext.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/HybridSearchDocumentQueryExecutionContext.java @@ -335,7 +335,7 @@ public Flux> apply(Flux } } - static Flux> computeRRFScores( + private static Flux> computeRRFScores( Mono>> ranks, Mono>> coalescedAndSortedResults, List componentWeights) { @@ -362,7 +362,7 @@ static Flux> computeRRFScores( }).flatMapMany(Flux::fromIterable); } - static Mono>> computeRanks( + private static Mono>> computeRanks( Mono>> componentScoresList, boolean useCompetitionRanking) { return componentScoresList.map(componentScores -> { // initialize ranks as an N-D list with zeros @@ -392,7 +392,7 @@ static Mono>> computeRanks( }); } - static Mono>> retrieveComponentScores( + private static Mono>> retrieveComponentScores( Mono>> coalescedAndSortedResults, List componentWeights) { return coalescedAndSortedResults.map(results -> { List> componentScoresInternal = new ArrayList<>(); @@ -519,7 +519,7 @@ private Flux retrieveRewrittenQueryInfos(List componentQue }); } - static Mono>> coalesceAndSortResults(Flux componentQueryResults) { + private static Mono>> coalesceAndSortResults(Flux componentQueryResults) { return componentQueryResults.collectList() .map(results -> { Map> uniqueDocuments = new LinkedHashMap<>(); @@ -591,7 +591,7 @@ private GlobalFullTextSearchQueryStatistics aggregateStatistics(List retrieveComponentWeights(List componentWeightList, List componentQueryInfos) { + private static List retrieveComponentWeights(List componentWeightList, List componentQueryInfos) { boolean useDefaultComponentWeight = componentWeightList == null || componentWeightList.isEmpty(); List componentWeights = new ArrayList<>(); for (int i=0;i retrieveComponentWeights(List componentWeig return componentWeights; } - static class ComponentWeight { + private static class ComponentWeight { private final Double weight; private final Comparator comparator; private final int comparisonFactor; From 959e3460f8f205dace710d5af16000a4c183ff76 Mon Sep 17 00:00:00 2001 From: Maya Painter Date: Mon, 26 Oct 2026 13:55:57 -0700 Subject: [PATCH 5/5] Consolidate and enable hybrid search integration tests Move ordinal and competition ranking coverage into HybridSearchQueryTest, reuse its database lifecycle, and remove the obsolete emulator ignore gate. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 29c3522f-9887-43c5-bebb-f42ee9ba4d51 --- .../cosmos/rx/HybridSearchQueryTest.java | 211 ++++++++++++++- .../cosmos/rx/HybridSearchRankingTest.java | 254 ------------------ 2 files changed, 209 insertions(+), 256 deletions(-) delete mode 100644 sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/HybridSearchRankingTest.java diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/HybridSearchQueryTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/HybridSearchQueryTest.java index c857412d02938..2569d32cb7c4b 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/HybridSearchQueryTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/HybridSearchQueryTest.java @@ -18,6 +18,7 @@ import com.azure.cosmos.models.CosmosFullTextPolicy; import com.azure.cosmos.models.CosmosFullTextScoreScope; import com.azure.cosmos.models.CosmosQueryRequestOptions; +import com.azure.cosmos.models.FeedResponse; import com.azure.cosmos.models.IncludedPath; import com.azure.cosmos.models.IndexingMode; import com.azure.cosmos.models.IndexingPolicy; @@ -32,19 +33,24 @@ import com.fasterxml.jackson.annotation.JsonProperty; import com.fasterxml.jackson.core.type.TypeReference; import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.JsonNodeFactory; +import com.fasterxml.jackson.databind.node.ObjectNode; import io.netty.handler.codec.http.HttpResponseStatus; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import org.testng.annotations.AfterClass; import org.testng.annotations.BeforeClass; -import org.testng.annotations.Ignore; +import org.testng.annotations.DataProvider; import org.testng.annotations.Test; import reactor.core.publisher.Flux; import java.io.File; import java.io.IOException; +import java.nio.ByteBuffer; +import java.nio.ByteOrder; import java.util.ArrayList; import java.util.Arrays; +import java.util.Base64; import java.util.Collections; import java.util.HashMap; import java.util.List; @@ -59,7 +65,7 @@ import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.fail; -@Ignore("TODO: Ignore these test cases until the public emulator is released.") +@Test(singleThreaded = true) public class HybridSearchQueryTest { protected static final int TIMEOUT = 30000; protected static final int SETUP_TIMEOUT = 80000; @@ -69,6 +75,8 @@ public class HybridSearchQueryTest { private static final List JOHN_AND_UNITED_IDS = Arrays.asList("2", "22", "24", "25", "49", "51", "54", "57", "61", "66", "75", "76", "77", "80", "85"); + private static final String RANKING_PROPERTY = "COSMOS.HYBRID_SEARCH_USE_COMPETITION_RANKING"; + private static final int RANKING_DOCUMENT_COUNT = 260; private final String containerId = UUID.randomUUID().toString(); private final String hpkContainerId = UUID.randomUUID().toString(); private CosmosAsyncClient client; @@ -80,6 +88,9 @@ public class HybridSearchQueryTest { private List textJohnOrder; private List textUnitedOrder; private List vectorOrder; + private CosmosAsyncContainer rankingContainer; + private CosmosAsyncContainer rankingTierContainer; + private List rankingDocumentsByRid; @BeforeClass(groups = {"query", "split"}, timeOut = SETUP_TIMEOUT* 10) public void before_HybridSearchQueryTest() { @@ -167,6 +178,7 @@ public void before_HybridSearchQueryTest() { return comparison != 0 ? comparison : Long.compareUnsigned(documentRids.get(right.id), documentRids.get(left.id)); }); vectorOrder = vectorDocuments.stream().map(Document::getId).collect(Collectors.toList()); + createRankingFixtures(); } @AfterClass(groups = {"query", "split"}, timeOut = SHUTDOWN_TIMEOUT, alwaysRun = true) @@ -660,6 +672,201 @@ public void wrongHybridQueryTest() { } + private void createRankingFixtures() { + CosmosFullTextIndex titleIndex = new CosmosFullTextIndex(); + titleIndex.setPath("/title"); + CosmosFullTextIndex textIndex = new CosmosFullTextIndex(); + textIndex.setPath("/text"); + IndexingPolicy indexing = new IndexingPolicy(); + indexing.setCosmosFullTextIndexes(Arrays.asList(titleIndex, textIndex)); + CosmosFullTextPath titlePath = new CosmosFullTextPath(); + titlePath.setPath("/title"); + titlePath.setLanguage("en-US"); + CosmosFullTextPath textPath = new CosmosFullTextPath(); + textPath.setPath("/text"); + textPath.setLanguage("en-US"); + CosmosFullTextPolicy fullText = new CosmosFullTextPolicy(); + fullText.setDefaultLanguage("en-US"); + fullText.setPaths(Arrays.asList(titlePath, textPath)); + CosmosContainerProperties properties = new CosmosContainerProperties("ranking-ties", "/pk"); + properties.setIndexingPolicy(indexing); + properties.setFullTextPolicy(fullText); + rankingContainer = createCollection(database, properties, new CosmosContainerRequestOptions(), 10000); + CosmosContainerProperties tierProperties = new CosmosContainerProperties("ranking-tiers", "/pk"); + tierProperties.setIndexingPolicy(indexing); + tierProperties.setFullTextPolicy(fullText); + rankingTierContainer = createCollection(database, tierProperties, new CosmosContainerRequestOptions(), 400); + String[] tierTexts = {"beta filler filler filler filler", "beta beta filler filler filler", + "beta beta beta filler filler", "beta beta beta beta filler", "beta beta beta beta beta"}; + for (int i = 0; i < 5; i++) { + ObjectNode document = JsonNodeFactory.instance.objectNode(); + document.put("id", "tier-" + i); + document.put("pk", "tiers"); + document.put("title", i == 0 ? "alpha alpha" : i == 4 ? "filler filler" : "alpha filler"); + document.put("text", tierTexts[i]); + rankingTierContainer.createItem(document).block(); + } + + // Identical fields guarantee tied component scores without reproducing the service's BM25 implementation. + // More than 256 inserts also exercise RID byte boundaries rather than assuming encoded RID order. + for (int i = 0; i < RANKING_DOCUMENT_COUNT; i++) { + ObjectNode document = JsonNodeFactory.instance.objectNode(); + document.put("id", "doc-" + i); + document.put("pk", Integer.toString(i % 2)); + document.put("title", "alpha"); + document.put("text", "beta"); + rankingContainer.createItem(document).block(); + } + rankingDocumentsByRid = rankingContainer.queryItems("SELECT * FROM c", + new CosmosQueryRequestOptions(), ObjectNode.class).collectList().block(); + assertThat(rankingDocumentsByRid).hasSize(RANKING_DOCUMENT_COUNT); + rankingDocumentsByRid.sort((left, right) -> Long.compareUnsigned(documentRid(right), documentRid(left))); + List encodedOrder = new ArrayList<>(rankingDocumentsByRid); + encodedOrder.sort((left, right) -> right.get("_rid").asText().compareTo(left.get("_rid").asText())); + assertThat(rankingIds(encodedOrder)).as("Fixture must distinguish numeric from encoded RID ordering") + .isNotEqualTo(rankingIds(rankingDocumentsByRid)); + } + + @DataProvider(name = "rankingModes") + public Object[][] rankingModes() { + return new Object[][] {{false}, {true}}; + } + + @Test(groups = {"query", "split"}, dataProvider = "rankingModes", timeOut = 60000) + public void tiedComponentsRespectRankingModeAndDirection(boolean competition) { + String previous = System.getProperty(RANKING_PROPERTY); + try { + System.setProperty(RANKING_PROPERTY, Boolean.toString(competition)); + for (double[] weights : new double[][] {{1, 1}, {-1, -1}, {1, -1}, {2, -1}}) { + List expected = expectedTiedRankingOrder(weights, competition); + String orderBy = rankingOrderBy(weights); + List> pages = rankingContainer.queryItems( + "SELECT TOP 260 * FROM c " + orderBy, new CosmosQueryRequestOptions(), ObjectNode.class) + .byPage(17).collectList().block(); + List actual = pages.stream().flatMap(page -> page.getResults().stream()) + .collect(Collectors.toList()); + assertThat(rankingIds(actual)).as("weights=%s, competition=%s", Arrays.toString(weights), competition) + .containsExactlyElementsOf(expected); + assertThat(pages).hasSizeGreaterThan(1); + assertThat(actual).allSatisfy(document -> { + assertThat(document.get("title").asText()).isEqualTo("alpha"); + assertThat(document.get("text").asText()).isEqualTo("beta"); + }); + + // Keep the full fixture in the service's bounded component candidate lists. + List window = rankingContainer.queryItems( + "SELECT * FROM c " + orderBy + " OFFSET 3 LIMIT 257", + new CosmosQueryRequestOptions(), ObjectNode.class).collectList().block(); + assertThat(rankingIds(window)).containsExactlyElementsOf(expected.subList(3, RANKING_DOCUMENT_COUNT)); + } + } finally { + restoreRankingProperty(previous); + } + } + + @Test(groups = {"query", "split"}, timeOut = TIMEOUT) + public void competitionRanksSkipTiedPositions() { + List documents = rankingTierContainer.queryItems("SELECT * FROM c", + new CosmosQueryRequestOptions(), ObjectNode.class).collectList().block(); + assertThat(documents).hasSize(5); + documents.sort((left, right) -> Long.compareUnsigned(documentRid(right), documentRid(left))); + List titleOrder = new ArrayList<>(documents); + // Equal document lengths: two occurrences outrank one, and no occurrence scores zero. + titleOrder.sort((left, right) -> Integer.compare(titleTier(left), titleTier(right))); + String previous = System.getProperty(RANKING_PROPERTY); + try { + for (boolean competition : new boolean[] {false, true}) { + System.setProperty(RANKING_PROPERTY, Boolean.toString(competition)); + List expected = new ArrayList<>(documents); + int[] titleRanks = competition ? new int[] {1, 2, 2, 2, 5} : new int[] {1, 2, 3, 4, 5}; + expected.sort((left, right) -> { + double leftScore = 1.0 / (60 + titleRanks[titleOrder.indexOf(left)]) + + 1.0 / (60 + textRank(left)); + double rightScore = 1.0 / (60 + titleRanks[titleOrder.indexOf(right)]) + + 1.0 / (60 + textRank(right)); + int comparison = Double.compare(rightScore, leftScore); + return comparison != 0 ? comparison : Long.compareUnsigned(documentRid(right), documentRid(left)); + }); + List results = rankingTierContainer.queryItems("SELECT TOP 5 * FROM c " + + rankingOrderBy(new double[] {1, 1}), new CosmosQueryRequestOptions(), ObjectNode.class) + .collectList().block(); + assertThat(rankingIds(results)).as("competition=%s", competition) + .containsExactlyElementsOf(rankingIds(expected)); + List topTwo = rankingTierContainer.queryItems("SELECT TOP 2 * FROM c " + + rankingOrderBy(new double[] {1, 1}), new CosmosQueryRequestOptions(), ObjectNode.class) + .collectList().block(); + assertThat(rankingIds(topTwo)).containsExactlyElementsOf(rankingIds(expected.subList(0, 2))); + } + } finally { + restoreRankingProperty(previous); + } + } + + @Test(groups = {"query", "split"}, timeOut = TIMEOUT) + public void emptyHybridQueryReturnsNoDocuments() { + List results = rankingContainer.queryItems( + "SELECT TOP 10 * FROM c WHERE c.id = 'missing' " + rankingOrderBy(new double[] {1, 1}), + new CosmosQueryRequestOptions(), ObjectNode.class).collectList().block(); + assertThat(results).isEmpty(); + } + + private List expectedTiedRankingOrder(double[] weights, boolean competition) { + List positions = new ArrayList<>(); + for (int i = 0; i < rankingDocumentsByRid.size(); i++) { + positions.add(i); + } + // In competition mode every tied score has rank 1. Ordinal rank follows the plan direction; + // negative weights reverse that direction before their absolute magnitude is used in fusion. + positions.sort((left, right) -> { + int comparison = Double.compare(tiedRankingScore(right, weights, competition), + tiedRankingScore(left, weights, competition)); + return comparison != 0 ? comparison : Integer.compare(left, right); + }); + return positions.stream().map(position -> rankingDocumentsByRid.get(position).get("id").asText()) + .collect(Collectors.toList()); + } + + private static double tiedRankingScore(int descendingPosition, double[] weights, boolean competition) { + double score = 0; + for (double weight : weights) { + int rank = competition ? 1 : weight < 0 ? RANKING_DOCUMENT_COUNT - descendingPosition : descendingPosition + 1; + score += Math.abs(weight) / (60 + rank); + } + return score; + } + + private static int titleTier(ObjectNode document) { + String title = document.get("title").asText(); + return "alpha alpha".equals(title) ? 0 : "alpha filler".equals(title) ? 1 : 2; + } + + private static int textRank(ObjectNode document) { + // Texts have equal lengths and one through five occurrences of beta, so all scores are distinct. + return 5 - Integer.parseInt(document.get("id").asText().substring("tier-".length())); + } + + private static String rankingOrderBy(double[] weights) { + return "ORDER BY RANK RRF(FullTextScore(c.title, 'alpha'), FullTextScore(c.text, 'beta'), " + + Arrays.toString(weights) + ")"; + } + + private static long documentRid(ObjectNode document) { + byte[] bytes = Base64.getDecoder().decode(document.get("_rid").asText().replace('-', '/')); + return ByteBuffer.wrap(bytes, 8, 8).order(ByteOrder.LITTLE_ENDIAN).getLong(); + } + + private static List rankingIds(List documents) { + return documents.stream().map(document -> document.get("id").asText()).collect(Collectors.toList()); + } + + private static void restoreRankingProperty(String previous) { + if (previous == null) { + System.clearProperty(RANKING_PROPERTY); + } else { + System.setProperty(RANKING_PROPERTY, previous); + } + } + private IndexingPolicy populateIndexingPolicy() { IndexingPolicy indexingPolicy = new IndexingPolicy(); indexingPolicy.setIndexingMode(IndexingMode.CONSISTENT); diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/HybridSearchRankingTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/HybridSearchRankingTest.java deleted file mode 100644 index 0e1bb6084600e..0000000000000 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/HybridSearchRankingTest.java +++ /dev/null @@ -1,254 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -package com.azure.cosmos.rx; - -import com.azure.cosmos.CosmosAsyncClient; -import com.azure.cosmos.CosmosAsyncContainer; -import com.azure.cosmos.CosmosAsyncDatabase; -import com.azure.cosmos.CosmosClientBuilder; -import com.azure.cosmos.implementation.TestConfigurations; -import com.azure.cosmos.models.CosmosContainerProperties; -import com.azure.cosmos.models.CosmosContainerRequestOptions; -import com.azure.cosmos.models.CosmosFullTextIndex; -import com.azure.cosmos.models.CosmosFullTextPath; -import com.azure.cosmos.models.CosmosFullTextPolicy; -import com.azure.cosmos.models.CosmosQueryRequestOptions; -import com.azure.cosmos.models.FeedResponse; -import com.azure.cosmos.models.IndexingPolicy; -import com.fasterxml.jackson.databind.node.JsonNodeFactory; -import com.fasterxml.jackson.databind.node.ObjectNode; -import org.testng.annotations.AfterClass; -import org.testng.annotations.BeforeClass; -import org.testng.annotations.DataProvider; -import org.testng.annotations.Test; - -import java.nio.ByteBuffer; -import java.nio.ByteOrder; -import java.util.ArrayList; -import java.util.Arrays; -import java.util.Base64; -import java.util.List; -import java.util.stream.Collectors; - -import static com.azure.cosmos.rx.TestSuiteBase.createCollection; -import static com.azure.cosmos.rx.TestSuiteBase.createTestDatabase; -import static com.azure.cosmos.rx.TestSuiteBase.safeClose; -import static com.azure.cosmos.rx.TestSuiteBase.safeDeleteDatabase; -import static org.assertj.core.api.Assertions.assertThat; - -// Runs in the query integration suite against an emulator/account with full-text search support. -@Test(groups = "query", singleThreaded = true) -public class HybridSearchRankingTest { - private static final String RANKING_PROPERTY = "COSMOS.HYBRID_SEARCH_USE_COMPETITION_RANKING"; - private static final int DOCUMENT_COUNT = 260; - private CosmosAsyncClient client; - private CosmosAsyncDatabase database; - private CosmosAsyncContainer container; - private CosmosAsyncContainer tierContainer; - private List documentsByRid; - - @BeforeClass(groups = "query", timeOut = 120000) - public void setUp() { - client = new CosmosClientBuilder().endpoint(TestConfigurations.HOST) - .key(TestConfigurations.MASTER_KEY).buildAsyncClient(); - database = createTestDatabase(client, "hybrid-ranking"); - - CosmosFullTextIndex titleIndex = new CosmosFullTextIndex(); - titleIndex.setPath("/title"); - CosmosFullTextIndex textIndex = new CosmosFullTextIndex(); - textIndex.setPath("/text"); - IndexingPolicy indexing = new IndexingPolicy(); - indexing.setCosmosFullTextIndexes(Arrays.asList(titleIndex, textIndex)); - CosmosFullTextPath titlePath = new CosmosFullTextPath(); - titlePath.setPath("/title"); - titlePath.setLanguage("en-US"); - CosmosFullTextPath textPath = new CosmosFullTextPath(); - textPath.setPath("/text"); - textPath.setLanguage("en-US"); - CosmosFullTextPolicy fullText = new CosmosFullTextPolicy(); - fullText.setDefaultLanguage("en-US"); - fullText.setPaths(Arrays.asList(titlePath, textPath)); - CosmosContainerProperties properties = new CosmosContainerProperties("ties", "/pk"); - properties.setIndexingPolicy(indexing); - properties.setFullTextPolicy(fullText); - container = createCollection(database, properties, new CosmosContainerRequestOptions(), 10000); - CosmosContainerProperties tierProperties = new CosmosContainerProperties("tiers", "/pk"); - tierProperties.setIndexingPolicy(indexing); - tierProperties.setFullTextPolicy(fullText); - tierContainer = createCollection(database, tierProperties, new CosmosContainerRequestOptions(), 400); - String[] tierTexts = {"beta filler filler filler filler", "beta beta filler filler filler", - "beta beta beta filler filler", "beta beta beta beta filler", "beta beta beta beta beta"}; - for (int i = 0; i < 5; i++) { - ObjectNode document = JsonNodeFactory.instance.objectNode(); - document.put("id", "tier-" + i); - document.put("pk", "tiers"); - document.put("title", i == 0 ? "alpha alpha" : i == 4 ? "filler filler" : "alpha filler"); - document.put("text", tierTexts[i]); - tierContainer.createItem(document).block(); - } - - // Identical fields guarantee tied component scores without reproducing the service's BM25 implementation. - // More than 256 inserts also exercise RID byte boundaries rather than assuming encoded RID order. - for (int i = 0; i < DOCUMENT_COUNT; i++) { - ObjectNode document = JsonNodeFactory.instance.objectNode(); - document.put("id", "doc-" + i); - document.put("pk", Integer.toString(i % 2)); - document.put("title", "alpha"); - document.put("text", "beta"); - container.createItem(document).block(); - } - documentsByRid = container.queryItems("SELECT * FROM c", new CosmosQueryRequestOptions(), ObjectNode.class) - .collectList().block(); - assertThat(documentsByRid).hasSize(DOCUMENT_COUNT); - documentsByRid.sort((left, right) -> Long.compareUnsigned(documentRid(right), documentRid(left))); - List encodedOrder = new ArrayList<>(documentsByRid); - encodedOrder.sort((left, right) -> right.get("_rid").asText().compareTo(left.get("_rid").asText())); - assertThat(ids(encodedOrder)).as("Fixture must distinguish numeric from encoded RID ordering") - .isNotEqualTo(ids(documentsByRid)); - } - - @AfterClass(groups = "query", alwaysRun = true, timeOut = 30000) - public void tearDown() { - safeDeleteDatabase(database); - safeClose(client); - } - - @DataProvider(name = "rankingModes") - public Object[][] rankingModes() { - return new Object[][] {{false}, {true}}; - } - - @Test(groups = "query", dataProvider = "rankingModes", timeOut = 60000) - public void tiedComponentsRespectRankingModeAndDirection(boolean competition) { - String previous = System.getProperty(RANKING_PROPERTY); - try { - System.setProperty(RANKING_PROPERTY, Boolean.toString(competition)); - for (double[] weights : new double[][] {{1, 1}, {-1, -1}, {1, -1}, {2, -1}}) { - List expected = expectedOrder(weights, competition); - String orderBy = orderBy(weights); - List> pages = container.queryItems( - "SELECT TOP 260 * FROM c " + orderBy, new CosmosQueryRequestOptions(), ObjectNode.class) - .byPage(17).collectList().block(); - List actual = pages.stream().flatMap(page -> page.getResults().stream()) - .collect(Collectors.toList()); - assertThat(ids(actual)).as("weights=%s, competition=%s", Arrays.toString(weights), competition) - .containsExactlyElementsOf(expected); - assertThat(pages).hasSizeGreaterThan(1); - assertThat(actual).allSatisfy(document -> { - assertThat(document.get("title").asText()).isEqualTo("alpha"); - assertThat(document.get("text").asText()).isEqualTo("beta"); - }); - - // Keep the full fixture in the service's bounded component candidate lists. - List window = container.queryItems( - "SELECT * FROM c " + orderBy + " OFFSET 3 LIMIT 257", - new CosmosQueryRequestOptions(), ObjectNode.class).collectList().block(); - assertThat(ids(window)).containsExactlyElementsOf(expected.subList(3, DOCUMENT_COUNT)); - } - } finally { - restoreRankingProperty(previous); - } - } - - @Test(groups = "query", timeOut = 30000) - public void competitionRanksSkipTiedPositions() { - List documents = tierContainer.queryItems("SELECT * FROM c", - new CosmosQueryRequestOptions(), ObjectNode.class).collectList().block(); - assertThat(documents).hasSize(5); - documents.sort((left, right) -> Long.compareUnsigned(documentRid(right), documentRid(left))); - List titleOrder = new ArrayList<>(documents); - // Equal document lengths: two occurrences outrank one, and no occurrence scores zero. - titleOrder.sort((left, right) -> Integer.compare(titleTier(left), titleTier(right))); - String previous = System.getProperty(RANKING_PROPERTY); - try { - for (boolean competition : new boolean[] {false, true}) { - System.setProperty(RANKING_PROPERTY, Boolean.toString(competition)); - List expected = new ArrayList<>(documents); - int[] titleRanks = competition ? new int[] {1, 2, 2, 2, 5} : new int[] {1, 2, 3, 4, 5}; - expected.sort((left, right) -> { - double leftScore = 1.0 / (60 + titleRanks[titleOrder.indexOf(left)]) - + 1.0 / (60 + textRank(left)); - double rightScore = 1.0 / (60 + titleRanks[titleOrder.indexOf(right)]) - + 1.0 / (60 + textRank(right)); - int comparison = Double.compare(rightScore, leftScore); - return comparison != 0 ? comparison : Long.compareUnsigned(documentRid(right), documentRid(left)); - }); - List results = tierContainer.queryItems("SELECT TOP 5 * FROM c " - + orderBy(new double[] {1, 1}), new CosmosQueryRequestOptions(), ObjectNode.class) - .collectList().block(); - assertThat(ids(results)).as("competition=%s", competition).containsExactlyElementsOf(ids(expected)); - List topTwo = tierContainer.queryItems("SELECT TOP 2 * FROM c " - + orderBy(new double[] {1, 1}), new CosmosQueryRequestOptions(), ObjectNode.class) - .collectList().block(); - assertThat(ids(topTwo)).containsExactlyElementsOf(ids(expected.subList(0, 2))); - } - } finally { - restoreRankingProperty(previous); - } - } - - private static int titleTier(ObjectNode document) { - String title = document.get("title").asText(); - return "alpha alpha".equals(title) ? 0 : "alpha filler".equals(title) ? 1 : 2; - } - - private static int textRank(ObjectNode document) { - // Texts have equal lengths and one through five occurrences of beta, so all scores are distinct. - return 5 - Integer.parseInt(document.get("id").asText().substring("tier-".length())); - } - - @Test(groups = "query", timeOut = 30000) - public void emptyHybridQueryReturnsNoDocuments() { - List results = container.queryItems( - "SELECT TOP 10 * FROM c WHERE c.id = 'missing' " + orderBy(new double[] {1, 1}), - new CosmosQueryRequestOptions(), ObjectNode.class).collectList().block(); - assertThat(results).isEmpty(); - } - - private List expectedOrder(double[] weights, boolean competition) { - List positions = new ArrayList<>(); - for (int i = 0; i < documentsByRid.size(); i++) { - positions.add(i); - } - // In competition mode every tied score has rank 1. Ordinal rank follows the plan direction; - // negative weights reverse that direction before their absolute magnitude is used in fusion. - positions.sort((left, right) -> { - int comparison = Double.compare(score(right, weights, competition), score(left, weights, competition)); - return comparison != 0 ? comparison : Integer.compare(left, right); - }); - return positions.stream().map(position -> documentsByRid.get(position).get("id").asText()) - .collect(Collectors.toList()); - } - - private static double score(int descendingPosition, double[] weights, boolean competition) { - double score = 0; - for (double weight : weights) { - int rank = competition ? 1 : weight < 0 ? DOCUMENT_COUNT - descendingPosition : descendingPosition + 1; - score += Math.abs(weight) / (60 + rank); - } - return score; - } - - private static String orderBy(double[] weights) { - return "ORDER BY RANK RRF(FullTextScore(c.title, 'alpha'), FullTextScore(c.text, 'beta'), " - + Arrays.toString(weights) + ")"; - } - - private static long documentRid(ObjectNode document) { - byte[] bytes = Base64.getDecoder().decode(document.get("_rid").asText().replace('-', '/')); - return ByteBuffer.wrap(bytes, 8, 8).order(ByteOrder.LITTLE_ENDIAN).getLong(); - } - - private static List ids(List documents) { - return documents.stream().map(document -> document.get("id").asText()).collect(Collectors.toList()); - } - - private static void restoreRankingProperty(String previous) { - if (previous == null) { - System.clearProperty(RANKING_PROPERTY); - } else { - System.setProperty(RANKING_PROPERTY, previous); - } - } -}