From 25ab6c14656f0a0ea49d1264a4c2f8c7e6578be5 Mon Sep 17 00:00:00 2001 From: Aman Rao Date: Mon, 5 May 2025 19:20:37 +0530 Subject: [PATCH 1/6] Add support for weighted RRF --- .../hybridQueryExecutionContext.ts | 93 ++++++++++++++++--- 1 file changed, 78 insertions(+), 15 deletions(-) diff --git a/sdk/cosmosdb/cosmos/src/queryExecutionContext/hybridQueryExecutionContext.ts b/sdk/cosmosdb/cosmos/src/queryExecutionContext/hybridQueryExecutionContext.ts index 545d1a114974..80bd4711b0e1 100644 --- a/sdk/cosmosdb/cosmos/src/queryExecutionContext/hybridQueryExecutionContext.ts +++ b/sdk/cosmosdb/cosmos/src/queryExecutionContext/hybridQueryExecutionContext.ts @@ -250,7 +250,11 @@ export class HybridQueryExecutionContext implements ExecutionContext { } // Initialize an array to hold ranks for each document - const sortedHybridSearchResult = this.sortHybridSearchResultByRRFScore(this.hybridSearchResult); + const componentWeights = this.extractComponentWeights(); + const sortedHybridSearchResult = this.sortHybridSearchResultByRRFScore( + this.hybridSearchResult, + componentWeights, + ); // store the result to buffer // add only data from the sortedHybridSearchResult in the buffer sortedHybridSearchResult.forEach((item) => this.buffer.push(item.data)); @@ -318,6 +322,7 @@ export class HybridQueryExecutionContext implements ExecutionContext { private sortHybridSearchResultByRRFScore( hybridSearchResult: HybridSearchQueryResult[], + componentWeights: ComponentWeight[], ): HybridSearchQueryResult[] { if (hybridSearchResult.length === 0) { return []; @@ -329,7 +334,9 @@ export class HybridQueryExecutionContext implements ExecutionContext { // Compute ranks for each component score for (let i = 0; i < hybridSearchResult[0].componentScores.length; i++) { // Sort based on the i-th component score - hybridSearchResult.sort((a, b) => b.componentScores[i] - a.componentScores[i]); + hybridSearchResult.sort((a, b) => + componentWeights[i].comparator(a.componentScores[i], b.componentScores[i]), + ); // Assign ranks let rank = 1; @@ -338,7 +345,7 @@ export class HybridQueryExecutionContext implements ExecutionContext { j > 0 && hybridSearchResult[j].componentScores[i] !== hybridSearchResult[j - 1].componentScores[i] ) { - rank = j + 1; + ++rank; } const rankIndex = ranksArray.findIndex( (rankItem) => rankItem.rid === hybridSearchResult[j].rid, @@ -347,20 +354,14 @@ export class HybridQueryExecutionContext implements ExecutionContext { } } - // Function to compute RRF score - const computeRRFScore = (ranks: number[], k: number): number => { - return ranks.reduce((acc, rank) => acc + 1 / (k + rank), 0); - }; - // Compute RRF scores and sort based on them const rrfScores = ranksArray.map((item) => ({ rid: item.rid, - rrfScore: computeRRFScore(item.ranks, this.RRF_CONSTANT), + rrfScore: this.computeRRFScore(item.ranks, this.RRF_CONSTANT, componentWeights), })); // Sort based on RRF scores rrfScores.sort((a, b) => b.rrfScore - a.rrfScore); - // Map sorted RRF scores back to hybridSearchResult const sortedHybridSearchResult = rrfScores.map((scoreItem) => hybridSearchResult.find((item) => item.rid === scoreItem.rid), @@ -455,8 +456,14 @@ export class HybridQueryExecutionContext implements ExecutionContext { globalStats: GlobalStatistics, ): QueryInfo[] { return componentQueryInfos.map((queryInfo) => { - if (!queryInfo.hasNonStreamingOrderBy) { - throw new Error("The component query must have a non-streaming order by clause."); + let rewrittenOrderByExpressions = queryInfo.orderByExpressions; + if (queryInfo.orderBy && queryInfo.orderBy.length > 0) { + if (!queryInfo.hasNonStreamingOrderBy) { + throw new Error("The component query must have a non-streaming order by clause."); + } + rewrittenOrderByExpressions = queryInfo.orderByExpressions.map((expr) => + this.replacePlaceholdersWorkaroud(expr, globalStats, componentQueryInfos.length), + ); } return { ...queryInfo, @@ -465,9 +472,7 @@ export class HybridQueryExecutionContext implements ExecutionContext { globalStats, componentQueryInfos.length, ), - orderByExpressions: queryInfo.orderByExpressions.map((expr) => - this.replacePlaceholdersWorkaroud(expr, globalStats, componentQueryInfos.length), - ), + orderByExpressions: rewrittenOrderByExpressions, }; }); } @@ -531,4 +536,62 @@ export class HybridQueryExecutionContext implements ExecutionContext { } return query; } + + private computeRRFScore = ( + ranks: number[], + k: number, + componentWeights: ComponentWeight[], + ): number => { + if (ranks.length !== componentWeights.length) { + throw new Error("Ranks and component weights length mismatch"); + } + let rrfScore = 0; + for (let i = 0; i < ranks.length; i++) { + const rank = ranks[i]; + const weight = componentWeights[i].weight; + rrfScore += weight * (1 / (k + rank)); + } + return rrfScore; + }; + + private extractComponentWeights(): ComponentWeight[] { + const hybridSearchQueryInfo = this.partitionedQueryExecutionInfo.hybridSearchQueryInfo; + const useDefaultComponentWeight = + !hybridSearchQueryInfo.componentWeights || + hybridSearchQueryInfo.componentWeights.length === 0; + + const result: { + weight: number; + comparator: (x: number, y: number) => number; + }[] = []; + + for (let index = 0; index < hybridSearchQueryInfo.componentQueryInfos.length; ++index) { + const queryInfo = hybridSearchQueryInfo.componentQueryInfos[index]; + + if (queryInfo.orderBy && queryInfo.orderBy.length > 0) { + if (!queryInfo.hasNonStreamingOrderBy) { + throw new Error("The component query should have a non streaming order by"); + } + + if (!queryInfo.orderByExpressions || queryInfo.orderByExpressions.length !== 1) { + throw new Error("The component query should have exactly one order by expression"); + } + } + const componentWeight = useDefaultComponentWeight + ? 1 + : hybridSearchQueryInfo.componentWeights[index]; + const hasOrderBy = queryInfo.orderBy && queryInfo.orderBy.length > 0; + const sortOrder = hasOrderBy && queryInfo.orderBy[0].includes("Ascending") ? 1 : -1; + result.push({ + weight: componentWeight, + comparator: (x: number, y: number) => sortOrder * (x - y), + }); + } + return result; + } +} + +export interface ComponentWeight { + weight: number; + comparator: (x: number, y: number) => number; } From 9d1a9008eb22296b09d8ebeb469ed05521eca673 Mon Sep 17 00:00:00 2001 From: Aman Rao Date: Mon, 5 May 2025 19:21:25 +0530 Subject: [PATCH 2/6] Add component weights --- sdk/cosmosdb/cosmos/src/request/ErrorResponse.ts | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/sdk/cosmosdb/cosmos/src/request/ErrorResponse.ts b/sdk/cosmosdb/cosmos/src/request/ErrorResponse.ts index 95011c330883..aac7c73d8fd3 100644 --- a/sdk/cosmosdb/cosmos/src/request/ErrorResponse.ts +++ b/sdk/cosmosdb/cosmos/src/request/ErrorResponse.ts @@ -80,6 +80,10 @@ export interface HybridSearchQueryInfo { * Whether the query requires global statistics */ requiresGlobalStatistics: boolean; + /** + * Represents the weights for each component in a hybrid search query. + */ + componentWeights?: number[]; } export type GroupByExpressions = string[]; From 953f83c49f106cf95fe7a05ad28e008a6f91a211 Mon Sep 17 00:00:00 2001 From: Aman Rao Date: Mon, 5 May 2025 19:21:58 +0530 Subject: [PATCH 3/6] Add query plan optimization --- .../src/request/hybridSearchQueryResult.ts | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/sdk/cosmosdb/cosmos/src/request/hybridSearchQueryResult.ts b/sdk/cosmosdb/cosmos/src/request/hybridSearchQueryResult.ts index a7751cfda98d..45d9f0ba8ede 100644 --- a/sdk/cosmosdb/cosmos/src/request/hybridSearchQueryResult.ts +++ b/sdk/cosmosdb/cosmos/src/request/hybridSearchQueryResult.ts @@ -29,20 +29,27 @@ export class HybridSearchQueryResult { } const outerPayload = document[FieldNames.Payload]; + let componentScores: number[]; + let data: Record; + if (!outerPayload || typeof outerPayload !== "object") { throw new Error(`${FieldNames.Payload} must exist.`); } const innerPayload = outerPayload[FieldNames.Payload]; - if (!innerPayload || typeof innerPayload !== "object") { - throw new Error(`${FieldNames.Payload} must exist nested within the outer payload field.`); - } - const componentScores = outerPayload[FieldNames.ComponentScores]; + if (innerPayload && typeof innerPayload === "object") { + // older format without query plan optimization + componentScores = outerPayload[FieldNames.ComponentScores]; + data = innerPayload; + } else { + // newer format with the optimization + componentScores = document[FieldNames.ComponentScores]; + data = outerPayload; + } if (!Array.isArray(componentScores)) { throw new Error(`${FieldNames.ComponentScores} must exist.`); } - - return new HybridSearchQueryResult(rid, componentScores, innerPayload); + return new HybridSearchQueryResult(rid, componentScores, data); } } From 8ed6599cb2e0667d93dee5fc8f3b5039101ff253 Mon Sep 17 00:00:00 2001 From: Aman Rao Date: Mon, 5 May 2025 19:22:37 +0530 Subject: [PATCH 4/6] Add Query feature and update builder --- sdk/cosmosdb/cosmos/src/ClientContext.ts | 4 +--- sdk/cosmosdb/cosmos/src/common/constants.ts | 2 ++ sdk/cosmosdb/cosmos/src/request/FeedOptions.ts | 6 +++++- .../src/utils/supportedQueryFeaturesBuilder.ts | 18 +++++++++++------- 4 files changed, 19 insertions(+), 11 deletions(-) diff --git a/sdk/cosmosdb/cosmos/src/ClientContext.ts b/sdk/cosmosdb/cosmos/src/ClientContext.ts index b2fcce98aed8..9bd0c5af9224 100644 --- a/sdk/cosmosdb/cosmos/src/ClientContext.ts +++ b/sdk/cosmosdb/cosmos/src/ClientContext.ts @@ -279,9 +279,7 @@ export class ClientContext { request.headers[HttpHeaders.IsQueryPlan] = "True"; request.headers[HttpHeaders.QueryVersion] = "1.4"; request.headers[HttpHeaders.ContentType] = QueryJsonContentType; - request.headers[HttpHeaders.SupportedQueryFeatures] = supportedQueryFeaturesBuilder( - options.disableNonStreamingOrderByQuery, - ); + request.headers[HttpHeaders.SupportedQueryFeatures] = supportedQueryFeaturesBuilder(options); if (typeof query === "string") { request.body = { query }; // Converts query text to query object. diff --git a/sdk/cosmosdb/cosmos/src/common/constants.ts b/sdk/cosmosdb/cosmos/src/common/constants.ts index a1caeffffcc1..bea23a6ad337 100644 --- a/sdk/cosmosdb/cosmos/src/common/constants.ts +++ b/sdk/cosmosdb/cosmos/src/common/constants.ts @@ -523,6 +523,8 @@ export enum QueryFeature { ListAndSetAggregate = "ListAndSetAggregate", CountIf = "CountIf", HybridSearch = "HybridSearch", + WeightedRankFusion = "WeightedRankFusion", + HybridSearchSkipOrderByRewrite = "HybridSearchSkipOrderByRewrite", } export enum SDKSupportedCapabilities { diff --git a/sdk/cosmosdb/cosmos/src/request/FeedOptions.ts b/sdk/cosmosdb/cosmos/src/request/FeedOptions.ts index 4f3fe7538b48..ab1f5b4d6786 100644 --- a/sdk/cosmosdb/cosmos/src/request/FeedOptions.ts +++ b/sdk/cosmosdb/cosmos/src/request/FeedOptions.ts @@ -128,13 +128,17 @@ export interface FeedOptions extends SharedOptions { * Default: false; When set to true, it allows queries to bypass the default behavior that blocks nonStreaming queries without top or limit clauses. */ allowUnboundedNonStreamingQueries?: boolean; - /** * Controls query execution behavior. * Default: false. If set to false, the query will retry until results are ready and `maxItemCount` is reached, which can take time for large partitions with relatively small data. * If set to true, scans partitions up to `maxDegreeOfParallelism`, adds results to the buffer, and returns what is available. If results are not ready, it returns an empty response. */ enableQueryControl?: boolean; + /** + * Default: false. If set to true, it disables the hybrid search query plan optimization. + * This optimization is enabled by default and is used to improve the performance of hybrid search queries. + */ + disableHybridSearchQueryPlanOptimization?: boolean; /** * @internal * rid of the container. diff --git a/sdk/cosmosdb/cosmos/src/utils/supportedQueryFeaturesBuilder.ts b/sdk/cosmosdb/cosmos/src/utils/supportedQueryFeaturesBuilder.ts index a67b690cf041..015283ad3c2b 100644 --- a/sdk/cosmosdb/cosmos/src/utils/supportedQueryFeaturesBuilder.ts +++ b/sdk/cosmosdb/cosmos/src/utils/supportedQueryFeaturesBuilder.ts @@ -2,13 +2,17 @@ // Licensed under the MIT License. import { QueryFeature } from "../common/index.js"; +import type { FeedOptions } from "../request/FeedOptions.js"; -export function supportedQueryFeaturesBuilder(disableNonStreamingOrderByQuery?: boolean): string { - if (disableNonStreamingOrderByQuery) { - return Object.keys(QueryFeature) - .filter((k) => k !== QueryFeature.NonStreamingOrderBy) - .join(", "); - } else { - return Object.keys(QueryFeature).join(", "); +export function supportedQueryFeaturesBuilder(options: FeedOptions): string { + const allFeatures = Object.keys(QueryFeature) as QueryFeature[]; + const exclude: QueryFeature[] = []; + + if (options.disableNonStreamingOrderByQuery) { + exclude.push(QueryFeature.NonStreamingOrderBy); + } + if (options.disableHybridSearchQueryPlanOptimization) { + exclude.push(QueryFeature.HybridSearchSkipOrderByRewrite); } + return allFeatures.filter((feature) => !exclude.includes(feature)).join(","); } From 32613e667f0bfaa7636f209c859268bad525bec3 Mon Sep 17 00:00:00 2001 From: Aman Rao Date: Mon, 5 May 2025 19:22:55 +0530 Subject: [PATCH 5/6] update test cases --- .../unit/hybridExecutionContext.spec.ts | 23 +++- .../supportedQueryFeaturesBuilder.spec.ts | 32 +++++- .../public/integration/fullTextSearch.spec.ts | 100 +++++++++++++++++- 3 files changed, 144 insertions(+), 11 deletions(-) diff --git a/sdk/cosmosdb/cosmos/test/internal/unit/hybridExecutionContext.spec.ts b/sdk/cosmosdb/cosmos/test/internal/unit/hybridExecutionContext.spec.ts index db37ccd8f83e..6b5af8daa91e 100644 --- a/sdk/cosmosdb/cosmos/test/internal/unit/hybridExecutionContext.spec.ts +++ b/sdk/cosmosdb/cosmos/test/internal/unit/hybridExecutionContext.spec.ts @@ -7,6 +7,7 @@ import { DiagnosticNodeType, } from "../../../src/index.js"; import type { ClientContext, FeedOptions, QueryInfo } from "../../../src/index.js"; +import type { ComponentWeight } from "../../../src/queryExecutionContext/hybridQueryExecutionContext.js"; import { HybridQueryExecutionContext, HybridQueryExecutionContextBaseStates, @@ -210,7 +211,12 @@ describe("hybridQueryExecutionContext", () => { ]; const expectedSortedRids = ["3", "2", "1"]; - const result = context["sortHybridSearchResultByRRFScore"](input); + const comparator = (x: number, y: number): number => -1 * (x - y); + const componentWeights = [ + { weight: 1, comparator }, + { weight: 1, comparator }, + ]; + const result = context["sortHybridSearchResultByRRFScore"](input, componentWeights); const resultRids = result.map((res) => res.rid); // Assert that the result rids are equal to the expected sorted rids assert.deepStrictEqual(resultRids, expectedSortedRids); @@ -222,9 +228,11 @@ describe("hybridQueryExecutionContext", () => { { rid: "2", componentScores: [20], data: {}, score: 0, ranks: [] }, { rid: "3", componentScores: [25], data: {}, score: 0, ranks: [] }, ]; + const comparator = (x: number, y: number): number => -1 * (x - y); + const componentWeights = [{ weight: 1, comparator }]; const expectedSortedRids = ["1", "3", "2"]; - const result = context["sortHybridSearchResultByRRFScore"](input); + const result = context["sortHybridSearchResultByRRFScore"](input, componentWeights); const resultRids = result.map((res) => res.rid); // Assert that the result rids are equal to the expected sorted rids assert.deepStrictEqual(resultRids, expectedSortedRids); @@ -236,7 +244,13 @@ describe("hybridQueryExecutionContext", () => { ]; const expectedSortedRids = ["1"]; - const result = context["sortHybridSearchResultByRRFScore"](input); + const comparator = (x: number, y: number): number => -1 * (x - y); + const componentWeights = [ + { weight: 1, comparator }, + { weight: 1, comparator }, + { weight: 1, comparator }, + ]; + const result = context["sortHybridSearchResultByRRFScore"](input, componentWeights); const resultRids = result.map((res) => res.rid); // Assert that the result rids are equal to the expected sorted rids assert.deepStrictEqual(resultRids, expectedSortedRids); @@ -244,7 +258,8 @@ describe("hybridQueryExecutionContext", () => { it("sortHybridSearchResultByRRFScore method should handle empty HybridSearchQueryResult array", async () => { const input: HybridSearchQueryResult[] = []; - const result = context["sortHybridSearchResultByRRFScore"](input); + const componentWeights: ComponentWeight[] = []; + const result = context["sortHybridSearchResultByRRFScore"](input, componentWeights); assert.deepStrictEqual(input, result); }); }); diff --git a/sdk/cosmosdb/cosmos/test/internal/unit/utils/supportedQueryFeaturesBuilder.spec.ts b/sdk/cosmosdb/cosmos/test/internal/unit/utils/supportedQueryFeaturesBuilder.spec.ts index 9464aa17fd99..b02b508a28f7 100644 --- a/sdk/cosmosdb/cosmos/test/internal/unit/utils/supportedQueryFeaturesBuilder.spec.ts +++ b/sdk/cosmosdb/cosmos/test/internal/unit/utils/supportedQueryFeaturesBuilder.spec.ts @@ -8,19 +8,43 @@ import { describe, it, assert } from "vitest"; describe("validate supportedQueryFeaturesBuilder", () => { it("should contain nonStreamingOrderBy feature", () => { const feedOptions: FeedOptions = {}; - const res = supportedQueryFeaturesBuilder(feedOptions.disableNonStreamingOrderByQuery); + const res = supportedQueryFeaturesBuilder(feedOptions); assert.equal(res.includes("NonStreamingOrderBy"), true); }); it("should contain nonStreamingOrderBy feature", () => { const feedOptions: FeedOptions = { disableNonStreamingOrderByQuery: false }; - const res = supportedQueryFeaturesBuilder(feedOptions.disableNonStreamingOrderByQuery); + const res = supportedQueryFeaturesBuilder(feedOptions); assert.equal(res.includes("NonStreamingOrderBy"), true); }); - it("should contain nonStreamingOrderBy feature", () => { + it("should not contain nonStreamingOrderBy feature", () => { const feedOptions: FeedOptions = { disableNonStreamingOrderByQuery: true }; - const res = supportedQueryFeaturesBuilder(feedOptions.disableNonStreamingOrderByQuery); + const res = supportedQueryFeaturesBuilder(feedOptions); + assert.equal(res.includes("NonStreamingOrderBy"), false); + }); + it("should contain hybridSearchSkipOrderByRewrite feature", () => { + const feedOptions: FeedOptions = {}; + const res = supportedQueryFeaturesBuilder(feedOptions); + assert.equal(res.includes("HybridSearchSkipOrderByRewrite"), true); + }); + it("should contain hybridSearchSkipOrderByRewrite feature", () => { + const feedOptions: FeedOptions = { disableHybridSearchQueryPlanOptimization: false }; + const res = supportedQueryFeaturesBuilder(feedOptions); + assert.equal(res.includes("HybridSearchSkipOrderByRewrite"), true); + }); + it("should not contain hybridSearchSkipOrderByRewrite feature", () => { + const feedOptions: FeedOptions = { disableHybridSearchQueryPlanOptimization: true }; + const res = supportedQueryFeaturesBuilder(feedOptions); + assert.equal(res.includes("HybridSearchSkipOrderByRewrite"), false); + }); + it("should not contain nonStreamingOrderBy and hybridSearchSkipOrderByRewrite features", () => { + const feedOptions: FeedOptions = { + disableNonStreamingOrderByQuery: true, + disableHybridSearchQueryPlanOptimization: true, + }; + const res = supportedQueryFeaturesBuilder(feedOptions); assert.equal(res.includes("NonStreamingOrderBy"), false); + assert.equal(res.includes("HybridSearchSkipOrderByRewrite"), false); }); }); diff --git a/sdk/cosmosdb/cosmos/test/public/integration/fullTextSearch.spec.ts b/sdk/cosmosdb/cosmos/test/public/integration/fullTextSearch.spec.ts index 1505797bc23d..40f7221e042c 100644 --- a/sdk/cosmosdb/cosmos/test/public/integration/fullTextSearch.spec.ts +++ b/sdk/cosmosdb/cosmos/test/public/integration/fullTextSearch.spec.ts @@ -9,7 +9,7 @@ import { } from "../common/TestHelpers.js"; import { describe, it, assert, beforeAll } from "vitest"; -describe("Validate full text search queries", { timeout: 20000 }, () => { +describe.skip("FTSQuery", { timeout: 20000 }, () => { const partitionKey = "id"; let container: Container; const containerDefinition: ContainerDefinition = { @@ -240,6 +240,45 @@ describe("Validate full text search queries", { timeout: 20000 }, () => { }, ], // TODO: Add test case of just RRF with vector search no FullTextScore + [ + `SELECT c.index AS Index, c.title AS Title, c.text AS Text + FROM c + WHERE FullTextContains(c.title, 'John') OR FullTextContains(c.text, 'John') OR FullTextContains(c.text, 'United States') + ORDER BY RANK RRF(FullTextScore(c.title, ['John']), FullTextScore(c.text, ['United States']), [1, 1])`, + { + expected1: [61, 51, 49, 54, 75, 24, 77, 76, 80, 25, 22, 2, 66, 57, 85], + expected2: [61, 51, 49, 54, 75, 24, 77, 76, 80, 25, 22, 2, 66, 85, 57], + }, + ], + [ + ` + SELECT c.index AS Index, c.title AS Title, c.text AS Text + FROM c + WHERE FullTextContains(c.title, 'John') OR FullTextContains(c.text, 'John') OR FullTextContains(c.text, 'United States') + ORDER BY RANK RRF(FullTextScore(c.title, ['John']), FullTextScore(c.text, ['United States']), [0.1, 0.1])`, + { + expected1: [61, 51, 49, 54, 75, 24, 77, 76, 80, 25, 22, 2, 66, 57, 85], + expected2: [61, 51, 49, 54, 75, 24, 77, 76, 80, 25, 22, 2, 66, 85, 57], + }, + ], + [ + `SELECT c.index AS Index, c.title AS Title, c.text AS Text FROM c WHERE FullTextContains(c.title, 'John') OR FullTextContains(c.text, 'John') OR FullTextContains(c.text, 'United States') + ORDER BY RANK RRF(FullTextScore(c.title, ['John']), FullTextScore(c.text, ['United States']), [10, 10])`, + { + expected1: [61, 51, 49, 54, 75, 24, 77, 76, 80, 25, 22, 2, 66, 57, 85], + expected2: [61, 51, 49, 54, 75, 24, 77, 76, 80, 25, 22, 2, 66, 85, 57], + }, + ], + [ + `SELECT c.index AS Index, c.title AS Title, c.text AS Text + FROM c + WHERE FullTextContains(c.title, 'John') OR FullTextContains(c.text, 'John') OR FullTextContains(c.text, 'United States') + ORDER BY RANK RRF(FullTextScore(c.title, ['John']), FullTextScore(c.text, ['United States']), [-1, -1])`, + { + expected1: [85, 57, 2, 66, 22, 25, 80, 76, 77, 24, 75, 54, 49, 51, 61], + expected2: [85, 57, 2, 66, 22, 25, 77, 76, 80, 75, 24, 49, 54, 51, 61], + }, + ], ]); const containerOptions = { offerThroughput: 25000 }; @@ -287,6 +326,58 @@ describe("Validate full text search queries", { timeout: 20000 }, () => { assert.ok(isMatch, `The indexes array did not match expected values for query:\n${query}`); } }); + it("FetchNext: should return correct expected values for all the queries with enableQueryControl", async () => { + for (const [query, { expected1, expected2 }] of queriesMap) { + const queryOptions = { + allowUnboundedNonStreamingQueries: true, + forceQueryPlan: true, + enableQueryControl: true, + }; + const queryIterator = container.items.query(query, queryOptions); + + const results: any[] = []; + while (queryIterator.hasMoreResults()) { + const { resources: result } = await queryIterator.fetchNext(); + if (result !== undefined) { + results.push(...result); + } + } + + const indexes = results.map((result) => result.Index); + const isMatch = + JSON.stringify(indexes) === JSON.stringify(expected1) || + JSON.stringify(indexes) === JSON.stringify(expected2); + + assert.ok(isMatch, `The indexes array did not match expected values for query:\n${query}`); + } + }); + + it("FetchNext: should return correct expected values for all the queries with enableQueryControl and hybrid query plan optimization disabled", async () => { + for (const [query, { expected1, expected2 }] of queriesMap) { + const queryOptions = { + allowUnboundedNonStreamingQueries: true, + forceQueryPlan: true, + enableQueryControl: true, + disableHybridSearchQueryPlanOptimization: true, + }; + const queryIterator = container.items.query(query, queryOptions); + + const results: any[] = []; + while (queryIterator.hasMoreResults()) { + const { resources: result } = await queryIterator.fetchNext(); + if (result !== undefined) { + results.push(...result); + } + } + + const indexes = results.map((result) => result.Index); + const isMatch = + JSON.stringify(indexes) === JSON.stringify(expected1) || + JSON.stringify(indexes) === JSON.stringify(expected2); + + assert.ok(isMatch, `The indexes array did not match expected values for query:\n${query}`); + } + }); it("FetchAll: should return correct expected values for all the queries", async () => { for (const [query, { expected1, expected2 }] of queriesMap) { @@ -304,9 +395,12 @@ describe("Validate full text search queries", { timeout: 20000 }, () => { } }); - it("FetchAll: should return correct expected values for all the queries", async () => { + it("FetchAll: should return correct expected values for all the queries with hybrid query plan optimization disabled", async () => { for (const [query, { expected1, expected2 }] of queriesMap) { - const queryOptions = { allowUnboundedNonStreamingQueries: true, enableQueryControl: true }; + const queryOptions = { + allowUnboundedNonStreamingQueries: true, + disableHybridSearchQueryPlanOptimization: true, + }; const queryIterator = container.items.query(query, queryOptions); const { resources: results } = await queryIterator.fetchAll(); From e26d0f6d31dc2f5228c7ea2863355feec45d5684 Mon Sep 17 00:00:00 2001 From: Aman Rao Date: Mon, 5 May 2025 22:29:18 +0530 Subject: [PATCH 6/6] update public contract --- sdk/cosmosdb/cosmos/review/cosmos.api.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/sdk/cosmosdb/cosmos/review/cosmos.api.md b/sdk/cosmosdb/cosmos/review/cosmos.api.md index 333a1bcde56a..72ff2fe108a5 100644 --- a/sdk/cosmosdb/cosmos/review/cosmos.api.md +++ b/sdk/cosmosdb/cosmos/review/cosmos.api.md @@ -1191,6 +1191,7 @@ export interface FeedOptions extends SharedOptions { continuation?: string; continuationToken?: string; continuationTokenLimitInKB?: number; + disableHybridSearchQueryPlanOptimization?: boolean; disableNonStreamingOrderByQuery?: boolean; enableQueryControl?: boolean; enableScanInQuery?: boolean; @@ -1321,6 +1322,7 @@ export enum HTTPMethod { // @public export interface HybridSearchQueryInfo { componentQueryInfos: QueryInfo[]; + componentWeights?: number[]; globalStatisticsQuery: string; requiresGlobalStatistics: boolean; skip: number;