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..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 @@ -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; @@ -17,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; @@ -31,32 +33,39 @@ 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; +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; 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; @@ -64,12 +73,24 @@ 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 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; 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; + private CosmosAsyncContainer rankingContainer; + private CosmosAsyncContainer rankingTierContainer; + private List rankingDocumentsByRid; @BeforeClass(groups = {"query", "split"}, timeOut = SETUP_TIMEOUT* 10) public void before_HybridSearchQueryTest() { @@ -79,7 +100,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 +150,35 @@ 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()); + createRankingFixtures(); } @AfterClass(groups = {"query", "split"}, timeOut = SHUTDOWN_TIMEOUT, alwaysRun = true) @@ -149,7 +199,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 +210,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 +223,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 +238,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 +253,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 +274,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 +293,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 +311,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 +330,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 +351,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 +384,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 +401,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 +555,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(); @@ -617,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); @@ -648,11 +898,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 +976,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 diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md index 5d735ccc34aba..c20348c6a4acf 100644 --- a/sdk/cosmos/azure-cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md @@ -9,6 +9,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 queries scoped to an effective partition key sub-range returning items outside the requested range during fan-out or partition splits, including `ORDER BY` queries and continuation resumption. - See [PR 50943](https://github.com/Azure/azure-sdk-for-java/pull/50943) * 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) 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..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 @@ -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) { + private 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) { + private 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) { + private 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) { + private 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) { + 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 componentWei private 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); }