From 8c8d3e955ca1069909901d9f1338f04f10d2dae7 Mon Sep 17 00:00:00 2001 From: neildsh <35383880+neildsh@users.noreply.github.com> Date: Tue, 11 Jun 2024 06:46:12 -0700 Subject: [PATCH] Query: Fixes plumbing VectorEmbeddingPolicy to ServiceInterop to choose correct default distance function (#4538) * Plumb the collection VectorEmbeddingPolicy to ServiceInterop * Add query plan baseline tests for vector search * Correct typo in the query for baseline test * Fix build errors * fix runtime issue in mock setup due to the extra argument for vector embedding policy --- .../CosmosQueryExecutionContextFactory.cs | 1 + .../QueryClient/ContainerQueryProperties.cs | 12 +- .../Core/QueryClient/CosmosQueryClient.cs | 3 +- .../Core/QueryPlan/QueryPartitionProvider.cs | 21 ++- .../Query/Core/QueryPlan/QueryPlanHandler.cs | 18 +- .../Core/QueryPlan/QueryPlanRetriever.cs | 6 +- .../DefaultDocumentQueryExecutionContext.cs | 3 +- .../DocumentQueryExecutionContextBase.cs | 6 +- .../DocumentQueryExecutionContextFactory.cs | 3 +- .../Query/v3Query/CosmosQueryClientCore.cs | 9 +- .../src/Routing/PartitionRoutingHelper.cs | 4 +- .../CosmosItemTests.cs | 5 +- .../FeedToken/QueryFeedTokenTests.cs | 5 +- .../Query/ParsingBenchmark.cs | 3 +- .../QueryPlanBaselineTests.VectorSearch.xml | 176 ++++++++++++++++++ .../Microsoft.Azure.Cosmos.Tests.csproj | 3 + ...misticDirectExecutionQueryBaselineTests.cs | 43 +++-- .../Query/Pipeline/FullPipelineTests.cs | 20 +- .../Query/QueryPartitionProviderTests.cs | 6 +- .../Query/QueryPlanBaselineTests.cs | 90 ++++++++- .../Query/QueryPlanRetrieverTests.cs | 18 +- .../Query/SubpartitionTests.cs | 22 ++- .../Routing/PartitionRoutingHelperTest.cs | 6 +- .../Tracing/TraceWriterBaselineTests.cs | 3 +- 24 files changed, 414 insertions(+), 72 deletions(-) create mode 100644 Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/QueryPlanBaselineTests.VectorSearch.xml diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs index ccc8168983..ceff8377b4 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs @@ -609,6 +609,7 @@ private static async Task GetPartitionedQueryExec inputParameters.SqlQuerySpec, cosmosQueryContext.ResourceTypeEnum, partitionKeyDefinition, + containerQueryProperties.VectorEmbeddingPolicy, inputParameters.PartitionKey != null, containerQueryProperties.GeospatialType, cosmosQueryContext.UseSystemPrefix, diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/QueryClient/ContainerQueryProperties.cs b/Microsoft.Azure.Cosmos/src/Query/Core/QueryClient/ContainerQueryProperties.cs index 279b7585b9..18d8c59a0b 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/QueryClient/ContainerQueryProperties.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/QueryClient/ContainerQueryProperties.cs @@ -13,12 +13,14 @@ internal readonly struct ContainerQueryProperties public ContainerQueryProperties( string resourceId, IReadOnlyList> effectivePartitionKeyRanges, - PartitionKeyDefinition partitionKeyDefinition, + PartitionKeyDefinition partitionKeyDefinition, + Cosmos.VectorEmbeddingPolicy vectorEmbeddingPolicy, Cosmos.GeospatialType geospatialType) { this.ResourceId = resourceId; this.EffectiveRangesForPartitionKey = effectivePartitionKeyRanges; - this.PartitionKeyDefinition = partitionKeyDefinition; + this.PartitionKeyDefinition = partitionKeyDefinition; + this.VectorEmbeddingPolicy = vectorEmbeddingPolicy; this.GeospatialType = geospatialType; } @@ -27,7 +29,11 @@ public ContainerQueryProperties( //A PartitionKey has one range when it is a full PartitionKey value. //It can span many it is a prefix PartitionKey for a sub-partitioned container. public IReadOnlyList> EffectiveRangesForPartitionKey { get; } - public PartitionKeyDefinition PartitionKeyDefinition { get; } + + public PartitionKeyDefinition PartitionKeyDefinition { get; } + + public Cosmos.VectorEmbeddingPolicy VectorEmbeddingPolicy { get; } + public Cosmos.GeospatialType GeospatialType { get; } } } \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/QueryClient/CosmosQueryClient.cs b/Microsoft.Azure.Cosmos/src/Query/Core/QueryClient/CosmosQueryClient.cs index f32ab1035c..5d8d3b1461 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/QueryClient/CosmosQueryClient.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/QueryClient/CosmosQueryClient.cs @@ -40,7 +40,8 @@ public abstract Task GetCachedContainerQueryProperties public abstract Task> TryGetPartitionedQueryExecutionInfoAsync( SqlQuerySpec sqlQuerySpec, Documents.ResourceType resourceType, - Documents.PartitionKeyDefinition partitionKeyDefinition, + Documents.PartitionKeyDefinition partitionKeyDefinition, + Cosmos.VectorEmbeddingPolicy vectorEmbeddingPolicy, bool requireFormattableOrderByQuery, bool isContinuationExpected, bool allowNonValueAggregateQuery, diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryPartitionProvider.cs b/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryPartitionProvider.cs index dcb17d6463..6bec60605e 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryPartitionProvider.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryPartitionProvider.cs @@ -120,6 +120,7 @@ public void Update(IDictionary queryengineConfiguration) public TryCatch TryGetPartitionedQueryExecutionInfo( string querySpecJsonString, PartitionKeyDefinition partitionKeyDefinition, + VectorEmbeddingPolicy vectorEmbeddingPolicy, bool requireFormattableOrderByQuery, bool isContinuationExpected, bool allowNonValueAggregateQuery, @@ -130,7 +131,8 @@ public TryCatch TryGetPartitionedQueryExecutionIn { TryCatch tryGetInternalQueryInfo = this.TryGetPartitionedQueryExecutionInfoInternal( querySpecJsonString: querySpecJsonString, - partitionKeyDefinition: partitionKeyDefinition, + partitionKeyDefinition: partitionKeyDefinition, + vectorEmbeddingPolicy: vectorEmbeddingPolicy, requireFormattableOrderByQuery: requireFormattableOrderByQuery, isContinuationExpected: isContinuationExpected, allowNonValueAggregateQuery: allowNonValueAggregateQuery, @@ -179,7 +181,8 @@ internal PartitionedQueryExecutionInfo ConvertPartitionedQueryExecutionInfo( internal TryCatch TryGetPartitionedQueryExecutionInfoInternal( string querySpecJsonString, - PartitionKeyDefinition partitionKeyDefinition, + PartitionKeyDefinition partitionKeyDefinition, + VectorEmbeddingPolicy vectorEmbeddingPolicy, bool requireFormattableOrderByQuery, bool isContinuationExpected, bool allowNonValueAggregateQuery, @@ -222,8 +225,12 @@ internal TryCatch TryGetPartitionedQueryE Span buffer = stackalloc byte[QueryPartitionProvider.InitialBufferSize]; uint errorCode; - uint serializedQueryExecutionInfoResultLength; - + uint serializedQueryExecutionInfoResultLength; + + string vectorEmbeddingPolicyString = vectorEmbeddingPolicy != null ? + JsonConvert.SerializeObject(vectorEmbeddingPolicy) : + null; + unsafe { ServiceInteropWrapper.PartitionKeyRangesApiOptions partitionKeyRangesApiOptions = @@ -241,13 +248,15 @@ internal TryCatch TryGetPartitionedQueryE fixed (byte* bytePtr = buffer) { - errorCode = ServiceInteropWrapper.GetPartitionKeyRangesFromQuery3( + errorCode = ServiceInteropWrapper.GetPartitionKeyRangesFromQuery4( this.serviceProvider, querySpecJsonString, partitionKeyRangesApiOptions, allParts, partsLengths, - (uint)partitionKeyDefinition.Paths.Count, + (uint)partitionKeyDefinition.Paths.Count, + vectorEmbeddingPolicyString, + vectorEmbeddingPolicyString?.Length ?? 0, new IntPtr(bytePtr), (uint)buffer.Length, out serializedQueryExecutionInfoResultLength); diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryPlanHandler.cs b/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryPlanHandler.cs index 30be0a4903..29de3bce8c 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryPlanHandler.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryPlanHandler.cs @@ -25,7 +25,8 @@ public QueryPlanHandler(CosmosQueryClient queryClient) public async Task> TryGetQueryPlanAsync( SqlQuerySpec sqlQuerySpec, Documents.ResourceType resourceType, - PartitionKeyDefinition partitionKeyDefinition, + PartitionKeyDefinition partitionKeyDefinition, + VectorEmbeddingPolicy vectorEmbeddingPolicy, QueryFeatures supportedQueryFeatures, bool hasLogicalPartitionKey, bool useSystemPrefix, @@ -47,7 +48,8 @@ public async Task> TryGetQueryPlanAsync( TryCatch tryGetQueryInfo = await this.TryGetQueryInfoAsync( sqlQuerySpec, resourceType, - partitionKeyDefinition, + partitionKeyDefinition, + vectorEmbeddingPolicy, hasLogicalPartitionKey, useSystemPrefix, geospatialType, @@ -75,7 +77,8 @@ public async Task> TryGetQueryPlanAsync( QueryFeatures supportedQueryFeatures, SqlQuerySpec sqlQuerySpec, Documents.ResourceType resourceType, - PartitionKeyDefinition partitionKeyDefinition, + PartitionKeyDefinition partitionKeyDefinition, + VectorEmbeddingPolicy vectorEmbeddingPolicy, bool hasLogicalPartitionKey, bool useSystemPrefix, GeospatialType geospatialType, @@ -96,7 +99,8 @@ public async Task> TryGetQueryPlanAsync( TryCatch tryGetQueryInfo = await this.TryGetQueryInfoAsync( sqlQuerySpec, resourceType, - partitionKeyDefinition, + partitionKeyDefinition, + vectorEmbeddingPolicy, hasLogicalPartitionKey, useSystemPrefix, geospatialType, @@ -115,7 +119,8 @@ public async Task> TryGetQueryPlanAsync( private Task> TryGetQueryInfoAsync( SqlQuerySpec sqlQuerySpec, Documents.ResourceType resourceType, - PartitionKeyDefinition partitionKeyDefinition, + PartitionKeyDefinition partitionKeyDefinition, + VectorEmbeddingPolicy vectorEmbeddingPolicy, bool hasLogicalPartitionKey, bool useSystemPrefix, Cosmos.GeospatialType geospatialType, @@ -126,7 +131,8 @@ private Task> TryGetQueryInfoAsync( return this.queryClient.TryGetPartitionedQueryExecutionInfoAsync( sqlQuerySpec: sqlQuerySpec, resourceType: resourceType, - partitionKeyDefinition: partitionKeyDefinition, + partitionKeyDefinition: partitionKeyDefinition, + vectorEmbeddingPolicy: vectorEmbeddingPolicy, requireFormattableOrderByQuery: true, isContinuationExpected: false, allowNonValueAggregateQuery: true, diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryPlanRetriever.cs b/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryPlanRetriever.cs index 5ac7f41043..67cee232dc 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryPlanRetriever.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryPlanRetriever.cs @@ -50,7 +50,8 @@ public static async Task GetQueryPlanWithServiceI CosmosQueryClient queryClient, SqlQuerySpec sqlQuerySpec, Documents.ResourceType resourceType, - PartitionKeyDefinition partitionKeyDefinition, + PartitionKeyDefinition partitionKeyDefinition, + VectorEmbeddingPolicy vectorEmbeddingPolicy, bool hasLogicalPartitionKey, GeospatialType geospatialType, bool useSystemPrefix, @@ -81,7 +82,8 @@ public static async Task GetQueryPlanWithServiceI TryCatch tryGetQueryPlan = await queryPlanHandler.TryGetQueryPlanAsync( sqlQuerySpec, resourceType, - partitionKeyDefinition, + partitionKeyDefinition, + vectorEmbeddingPolicy, QueryPlanRetriever.SupportedQueryFeatures, hasLogicalPartitionKey, useSystemPrefix, diff --git a/Microsoft.Azure.Cosmos/src/Query/v2Query/DefaultDocumentQueryExecutionContext.cs b/Microsoft.Azure.Cosmos/src/Query/v2Query/DefaultDocumentQueryExecutionContext.cs index 9895bd5f5a..fa20ce39a3 100644 --- a/Microsoft.Azure.Cosmos/src/Query/v2Query/DefaultDocumentQueryExecutionContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/v2Query/DefaultDocumentQueryExecutionContext.cs @@ -316,7 +316,8 @@ private static bool ServiceInteropAvailable() allowDCount: false, allowNonValueAggregates: false, useSystemPrefix: false, - partitionKeyDefinition: partitionKeyDefinition, + partitionKeyDefinition: partitionKeyDefinition, + vectorEmbeddingPolicy: collection.VectorEmbeddingPolicy, queryPartitionProvider: queryPartitionProvider, clientApiVersion: version, geospatialType: collection.GeospatialConfig.GeospatialType, diff --git a/Microsoft.Azure.Cosmos/src/Query/v2Query/DocumentQueryExecutionContextBase.cs b/Microsoft.Azure.Cosmos/src/Query/v2Query/DocumentQueryExecutionContextBase.cs index 5bd47c8a21..bbcabfc41c 100644 --- a/Microsoft.Azure.Cosmos/src/Query/v2Query/DocumentQueryExecutionContextBase.cs +++ b/Microsoft.Azure.Cosmos/src/Query/v2Query/DocumentQueryExecutionContextBase.cs @@ -164,7 +164,8 @@ protected SqlQuerySpec QuerySpec public Guid CorrelatedActivityId { get; } public async Task GetPartitionedQueryExecutionInfoAsync( - PartitionKeyDefinition partitionKeyDefinition, + PartitionKeyDefinition partitionKeyDefinition, + Cosmos.VectorEmbeddingPolicy vectorEmbeddingPolicy, bool requireFormattableOrderByQuery, bool isContinuationExpected, bool allowNonValueAggregateQuery, @@ -179,7 +180,8 @@ public async Task GetPartitionedQueryExecutionInf QueryPartitionProvider queryPartitionProvider = await this.Client.GetQueryPartitionProviderAsync(); TryCatch tryGetPartitionedQueryExecutionInfo = queryPartitionProvider.TryGetPartitionedQueryExecutionInfo( querySpecJsonString: JsonConvert.SerializeObject(this.QuerySpec), - partitionKeyDefinition: partitionKeyDefinition, + partitionKeyDefinition: partitionKeyDefinition, + vectorEmbeddingPolicy: vectorEmbeddingPolicy, requireFormattableOrderByQuery: requireFormattableOrderByQuery, isContinuationExpected: isContinuationExpected, allowNonValueAggregateQuery: allowNonValueAggregateQuery, diff --git a/Microsoft.Azure.Cosmos/src/Query/v2Query/DocumentQueryExecutionContextFactory.cs b/Microsoft.Azure.Cosmos/src/Query/v2Query/DocumentQueryExecutionContextFactory.cs index f7861f5a2c..f167dfd84e 100644 --- a/Microsoft.Azure.Cosmos/src/Query/v2Query/DocumentQueryExecutionContextFactory.cs +++ b/Microsoft.Azure.Cosmos/src/Query/v2Query/DocumentQueryExecutionContextFactory.cs @@ -98,7 +98,8 @@ public static async Task CreateDocumentQueryExec //if collection is deleted/created with same name. //need to make it not rely on information from collection cache. PartitionedQueryExecutionInfo partitionedQueryExecutionInfo = await queryExecutionContext.GetPartitionedQueryExecutionInfoAsync( - partitionKeyDefinition: collection.PartitionKey, + partitionKeyDefinition: collection.PartitionKey, + vectorEmbeddingPolicy: collection.VectorEmbeddingPolicy, requireFormattableOrderByQuery: true, isContinuationExpected: isContinuationExpected, allowNonValueAggregateQuery: true, diff --git a/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs b/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs index f4817997ab..cd8f43ce19 100644 --- a/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs +++ b/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs @@ -85,14 +85,16 @@ public override async Task GetCachedContainerQueryProp return new ContainerQueryProperties( containerProperties.ResourceId, effectivePartitionKeyRange, - containerProperties.PartitionKey, + containerProperties.PartitionKey, + containerProperties.VectorEmbeddingPolicy, containerProperties.GeospatialConfig.GeospatialType); } public override async Task> TryGetPartitionedQueryExecutionInfoAsync( SqlQuerySpec sqlQuerySpec, ResourceType resourceType, - PartitionKeyDefinition partitionKeyDefinition, + PartitionKeyDefinition partitionKeyDefinition, + VectorEmbeddingPolicy vectorEmbeddingPolicy, bool requireFormattableOrderByQuery, bool isContinuationExpected, bool allowNonValueAggregateQuery, @@ -116,7 +118,8 @@ public override async Task> TryGetPartit return (await this.documentClient.QueryPartitionProvider).TryGetPartitionedQueryExecutionInfo( querySpecJsonString: queryString, - partitionKeyDefinition: partitionKeyDefinition, + partitionKeyDefinition: partitionKeyDefinition, + vectorEmbeddingPolicy: vectorEmbeddingPolicy, requireFormattableOrderByQuery: requireFormattableOrderByQuery, isContinuationExpected: isContinuationExpected, allowNonValueAggregateQuery: allowNonValueAggregateQuery, diff --git a/Microsoft.Azure.Cosmos/src/Routing/PartitionRoutingHelper.cs b/Microsoft.Azure.Cosmos/src/Routing/PartitionRoutingHelper.cs index 68f6214f17..6077aea787 100644 --- a/Microsoft.Azure.Cosmos/src/Routing/PartitionRoutingHelper.cs +++ b/Microsoft.Azure.Cosmos/src/Routing/PartitionRoutingHelper.cs @@ -35,6 +35,7 @@ public static IReadOnlyList> GetProvidedPartitionKeyRanges( bool allowNonValueAggregates, bool useSystemPrefix, PartitionKeyDefinition partitionKeyDefinition, + Cosmos.VectorEmbeddingPolicy vectorEmbeddingPolicy, QueryPartitionProvider queryPartitionProvider, string clientApiVersion, Cosmos.GeospatialType geospatialType, @@ -57,7 +58,8 @@ public static IReadOnlyList> GetProvidedPartitionKeyRanges( TryCatch tryGetPartitionQueryExecutionInfo = queryPartitionProvider.TryGetPartitionedQueryExecutionInfo( querySpecJsonString: querySpecJsonString, - partitionKeyDefinition: partitionKeyDefinition, + partitionKeyDefinition: partitionKeyDefinition, + vectorEmbeddingPolicy: vectorEmbeddingPolicy, requireFormattableOrderByQuery: VersionUtility.IsLaterThan(clientApiVersion, HttpConstants.VersionDates.v2016_11_14), isContinuationExpected: isContinuationExpected, allowNonValueAggregateQuery: allowNonValueAggregates, diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemTests.cs index 1463321ee7..0ad71e7085 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemTests.cs @@ -1723,9 +1723,10 @@ public async Task ItemEpkQuerySingleKeyRangeValidation() ContainerQueryProperties containerQueryProperties = new ContainerQueryProperties( containerResponse.Resource.ResourceId, - null, + effectivePartitionKeyRanges: null, //new List> { new Documents.Routing.Range("AA", "AA", true, true) }, - containerResponse.Resource.PartitionKey, + containerResponse.Resource.PartitionKey, + vectorEmbeddingPolicy: null, containerResponse.Resource.GeospatialConfig.GeospatialType); // There should only be one range since the EPK option is set. diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/FeedToken/QueryFeedTokenTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/FeedToken/QueryFeedTokenTests.cs index a5652e1010..74a8cf6f7c 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/FeedToken/QueryFeedTokenTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/FeedToken/QueryFeedTokenTests.cs @@ -54,8 +54,9 @@ public async Task GetTargetPartitionKeyRangesAsyncWithFeedRange() ContainerQueryProperties containerQueryProperties = new ContainerQueryProperties( containerResponse.Resource.ResourceId, - null, - containerResponse.Resource.PartitionKey, + effectivePartitionKeyRanges: null, + containerResponse.Resource.PartitionKey, + vectorEmbeddingPolicy: null, containerResponse.Resource.GeospatialConfig.GeospatialType); IReadOnlyList feedTokens = await container.GetFeedRangesAsync(); diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Performance.Tests/Query/ParsingBenchmark.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Performance.Tests/Query/ParsingBenchmark.cs index 2a17267604..62363f8a58 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Performance.Tests/Query/ParsingBenchmark.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Performance.Tests/Query/ParsingBenchmark.cs @@ -120,7 +120,8 @@ private static void ParseUsingNativeParser(SqlQuerySpec sqlQuerySpec) { TryCatch tryGetQueryPlan = QueryPartitionProvider.TryGetPartitionedQueryExecutionInfo( querySpecJsonString: JsonConvert.SerializeObject(sqlQuerySpec), - partitionKeyDefinition: PartitionKeyDefinition, + partitionKeyDefinition: PartitionKeyDefinition, + vectorEmbeddingPolicy: null, requireFormattableOrderByQuery: true, isContinuationExpected: false, allowNonValueAggregateQuery: true, diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/QueryPlanBaselineTests.VectorSearch.xml b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/QueryPlanBaselineTests.VectorSearch.xml new file mode 100644 index 0000000000..0b3074ec9d --- /dev/null +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/QueryPlanBaselineTests.VectorSearch.xml @@ -0,0 +1,176 @@ + + + + Euclidean Distance + SELECT TOP 10 c.title AS Title, VectorDistance(c.embedding, @vectorEmbedding, true) AS SimilarityScore + FROM c + ORDER BY VectorDistance(c.embedding, @vectorEmbedding, true) + + /PartitionKey + + Hash + Geography + + + + + + None + 10 + + + + + Ascending + + + VectorDistance(c.embedding, @vectorEmbedding, true) + + + + + False + + + + [[],"Infinity") + + + + + + + + + Cosine Similarity + SELECT TOP 10 c.title AS Title, VectorDistance(c.embedding, @vectorEmbedding, true) AS SimilarityScore + FROM c + ORDER BY VectorDistance(c.embedding, @vectorEmbedding, true) + + /PartitionKey + + Hash + Geography + + + + + + None + 10 + + + + + Descending + + + VectorDistance(c.embedding, @vectorEmbedding, true) + + + + + False + + + + [[],"Infinity") + + + + + + + + + Dot Product + SELECT TOP 10 c.title AS Title, VectorDistance(c.embedding, @vectorEmbedding, true) AS SimilarityScore + FROM c + ORDER BY VectorDistance(c.embedding, @vectorEmbedding, true) + + /PartitionKey + + Hash + Geography + + + + + + None + 10 + + + + + Descending + + + VectorDistance(c.embedding, @vectorEmbedding, true) + + + + + False + + + + [[],"Infinity") + + + + + + + \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Microsoft.Azure.Cosmos.Tests.csproj b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Microsoft.Azure.Cosmos.Tests.csproj index c9ad45d555..0262d1833d 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Microsoft.Azure.Cosmos.Tests.csproj +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Microsoft.Azure.Cosmos.Tests.csproj @@ -260,6 +260,9 @@ PreserveNewest + + PreserveNewest + PreserveNewest diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs index 03f6807e1d..267ea76b0e 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs @@ -30,7 +30,6 @@ using Microsoft.Azure.Documents; using Microsoft.Azure.Documents.Routing; using Microsoft.VisualStudio.TestTools.UnitTesting; - using Moq; using Newtonsoft.Json; using Newtonsoft.Json.Linq; @@ -143,7 +142,7 @@ await this.GetPipelineAndDrainAsync( clientDisableOde: clientForceDisableODEFromBackend); Assert.Fail("Expected exception. Received none."); } - catch(InvalidOperationException ex) + catch (InvalidOperationException ex) { Assert.IsTrue(ex.ToString().Contains("Execution of this query cannot resume using Optimistic Direct Execution continuation token due to partition split. Please restart the query without the continuation token.")); } @@ -604,8 +603,8 @@ public async Task TestClientDisableOdeLogic() public void TestTextDistributionPlanParsingFromStream() { string textPath = "../../../Query/DistributionPlans/Text"; - string[] filePaths = Directory.GetFiles(textPath); - + string[] filePaths = Directory.GetFiles(textPath); + foreach (string filePath in filePaths) { string testResponse = File.ReadAllText(filePath); @@ -820,15 +819,16 @@ internal static Tuple Get QueryPartitionProvider queryPartitionProvider = CreateCustomQueryPartitionProvider("clientDisableOptimisticDirectExecution", clientDisableOde.ToString().ToLower()); TryCatch tryGetQueryPlan = queryPartitionProvider.TryGetPartitionedQueryExecutionInfo( querySpecJsonString: querySpecJsonString, - partitionKeyDefinition: pkDefinition, + partitionKeyDefinition: pkDefinition, + vectorEmbeddingPolicy: null, requireFormattableOrderByQuery: true, isContinuationExpected: true, allowNonValueAggregateQuery: true, hasLogicalPartitionKey: false, allowDCount: true, useSystemPrefix: false, - geospatialType: Cosmos.GeospatialType.Geography); - + geospatialType: Cosmos.GeospatialType.Geography); + PartitionedQueryExecutionInfo partitionedQueryExecutionInfo = tryGetQueryPlan.Succeeded ? tryGetQueryPlan.Result : throw tryGetQueryPlan.Exception; return Tuple.Create(partitionedQueryExecutionInfo, queryPartitionProvider); } @@ -1025,8 +1025,8 @@ public override OptimisticDirectExecutionTestOutput ExecuteTest(OptimisticDirect isNonStreamingOrderByQueryFeatureDisabled: false, testInjections: queryRequestOptions.TestSettings); - List targetPkRanges = new (); - foreach(FeedRangeEpk feedRangeEpk in containerRanges) + List targetPkRanges = new(); + foreach (FeedRangeEpk feedRangeEpk in containerRanges) { targetPkRanges.Add(new PartitionKeyRange { @@ -1266,7 +1266,7 @@ internal class TestCosmosQueryClient : CosmosQueryClient private readonly IReadOnlyList targetPartitionKeyRanges; public TestCosmosQueryClient(QueryPartitionProvider queryPartitionProvider, IEnumerable targetPartitionKeyRanges) - { + { this.queryPartitionProvider = queryPartitionProvider; this.targetPartitionKeyRanges = targetPartitionKeyRanges.ToList(); } @@ -1303,14 +1303,15 @@ public override Task GetCachedContainerQueryProperties return Task.FromResult(new ContainerQueryProperties( "test", new List> - { + { new Range( PartitionKeyInternal.MinimumInclusiveEffectivePartitionKey, PartitionKeyInternal.MaximumExclusiveEffectivePartitionKey, true, true) }, - new PartitionKeyDefinition(), + new PartitionKeyDefinition(), + vectorEmbeddingPolicy: null, Cosmos.GeospatialType.Geometry)); } @@ -1334,12 +1335,24 @@ public override Task> TryGetOverlappingRangesAs throw new NotImplementedException(); } - public override Task> TryGetPartitionedQueryExecutionInfoAsync(SqlQuerySpec sqlQuerySpec, ResourceType resourceType, PartitionKeyDefinition partitionKeyDefinition, bool requireFormattableOrderByQuery, bool isContinuationExpected, bool allowNonValueAggregateQuery, bool hasLogicalPartitionKey, bool allowDCount, bool useSystemPrefix, Cosmos.GeospatialType geospatialType, CancellationToken cancellationToken) + public override Task> TryGetPartitionedQueryExecutionInfoAsync( + SqlQuerySpec sqlQuerySpec, + ResourceType resourceType, + PartitionKeyDefinition partitionKeyDefinition, + Cosmos.VectorEmbeddingPolicy vectorEmbeddingPolicy, + bool requireFormattableOrderByQuery, + bool isContinuationExpected, + bool allowNonValueAggregateQuery, + bool hasLogicalPartitionKey, + bool allowDCount, + bool useSystemPrefix, + Cosmos.GeospatialType geospatialType, + CancellationToken cancellationToken) { CosmosSerializerCore serializerCore = new CosmosSerializerCore(); using StreamReader streamReader = new StreamReader(serializerCore.ToStreamSqlQuerySpec(sqlQuerySpec, Documents.ResourceType.Document)); - string sqlQuerySpecJsonString = streamReader.ReadToEnd(); - + string sqlQuerySpecJsonString = streamReader.ReadToEnd(); + (PartitionedQueryExecutionInfo partitionedQueryExecutionInfo, QueryPartitionProvider queryPartitionProvider) = OptimisticDirectExecutionQueryBaselineTests.GetPartitionedQueryExecutionInfoAndPartitionProvider(sqlQuerySpecJsonString, partitionKeyDefinition); return Task.FromResult(TryCatch.FromResult(partitionedQueryExecutionInfo)); } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/FullPipelineTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/FullPipelineTests.cs index 298387bbff..3e56b607bd 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/FullPipelineTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/FullPipelineTests.cs @@ -375,7 +375,8 @@ private async Task TestPageSizeAsync(string query, int expectedPageSize, int exp mockClient.Setup(x => x.TryGetPartitionedQueryExecutionInfoAsync( It.IsAny(), It.IsAny(), - It.IsAny(), + It.IsAny(), + It.IsAny(), It.IsAny(), It.IsAny(), It.IsAny(), @@ -384,7 +385,19 @@ private async Task TestPageSizeAsync(string query, int expectedPageSize, int exp It.IsAny(), It.IsAny(), It.IsAny())) - .Returns((SqlQuerySpec sqlQuerySpec, ResourceType resourceType, PartitionKeyDefinition partitionKeyDefinition, bool requireFormattableOrderByQuery, bool isContinuationExpected, bool allowNonValueAggregateQuery, bool hasLogicalPartitionKey, bool allowDCount, bool useSystemPrefix, Cosmos.GeospatialType geospatialType, CancellationToken cancellationToken) => + .Returns(( + SqlQuerySpec sqlQuerySpec, + ResourceType resourceType, + PartitionKeyDefinition partitionKeyDefinition, + VectorEmbeddingPolicy vectorEmbeddingPolicy, + bool requireFormattableOrderByQuery, + bool isContinuationExpected, + bool allowNonValueAggregateQuery, + bool hasLogicalPartitionKey, + bool allowDCount, + bool useSystemPrefix, + Cosmos.GeospatialType geospatialType, + CancellationToken cancellationToken) => { CosmosSerializerCore serializerCore = new(); using StreamReader streamReader = new(serializerCore.ToStreamSqlQuerySpec(sqlQuerySpec, Documents.ResourceType.Document)); @@ -611,7 +624,8 @@ private static QueryInfo GetQueryPlan(string query) { TryCatch info = QueryPartitionProviderTestInstance.Object.TryGetPartitionedQueryExecutionInfoInternal( JsonConvert.SerializeObject(new SqlQuerySpec(query)), - partitionKeyDefinition, + partitionKeyDefinition, + vectorEmbeddingPolicy: null, requireFormattableOrderByQuery: true, isContinuationExpected: false, allowNonValueAggregateQuery: true, diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPartitionProviderTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPartitionProviderTests.cs index e4e6f48846..ff09f32052 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPartitionProviderTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPartitionProviderTests.cs @@ -33,7 +33,8 @@ public void TestQueryPartitionProviderUpdate() TryCatch tryGetQueryPlan = queryPartitionProvider.TryGetPartitionedQueryExecutionInfo( querySpecJsonString: sqlQuerySpec, - partitionKeyDefinition: PartitionKeyDefinition, + partitionKeyDefinition: PartitionKeyDefinition, + vectorEmbeddingPolicy: null, requireFormattableOrderByQuery: true, isContinuationExpected: false, allowNonValueAggregateQuery: true, @@ -49,7 +50,8 @@ public void TestQueryPartitionProviderUpdate() tryGetQueryPlan = queryPartitionProvider.TryGetPartitionedQueryExecutionInfo( querySpecJsonString: sqlQuerySpec, - partitionKeyDefinition: PartitionKeyDefinition, + partitionKeyDefinition: PartitionKeyDefinition, + vectorEmbeddingPolicy: null, requireFormattableOrderByQuery: true, isContinuationExpected: false, allowNonValueAggregateQuery: true, diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPlanBaselineTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPlanBaselineTests.cs index 204d3608d0..0717a80e33 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPlanBaselineTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPlanBaselineTests.cs @@ -1342,14 +1342,69 @@ public void Spatial() PartitionKeyDefinition pkDefinitions = CreateHashPartitionKey("/key"); return new List { - new QueryPlanBaselineTestInput($"{variation.Description} Geography", pkDefinitions, new SqlQuerySpec(variation.Query)) { GeospatialType = Cosmos.GeospatialType.Geography }, - new QueryPlanBaselineTestInput($"{variation.Description} Geometry", pkDefinitions, new SqlQuerySpec(variation.Query)) { GeospatialType = Cosmos.GeospatialType.Geometry } + new QueryPlanBaselineTestInput($"{variation.Description} Geography", pkDefinitions, vectorEmbeddingPolicy: null, new SqlQuerySpec(variation.Query), Cosmos.GeospatialType.Geography), + new QueryPlanBaselineTestInput($"{variation.Description} Geometry", pkDefinitions, vectorEmbeddingPolicy : null, new SqlQuerySpec(variation.Query), Cosmos.GeospatialType.Geometry) }; }) .ToList(); this.ExecuteTestSuite(testVariations); - } + } + + [TestMethod] + [Owner("ndeshpan")] + public void VectorSearch() + { + List testCases = new List + { + MakeVectorTest("Euclidean Distance", Cosmos.DistanceFunction.Euclidean), + MakeVectorTest("Cosine Similarity", Cosmos.DistanceFunction.Cosine), + MakeVectorTest("Dot Product", Cosmos.DistanceFunction.DotProduct), + }; + + this.ExecuteTestSuite(testCases); + } + + private static QueryPlanBaselineTestInput MakeVectorTest(string description, Cosmos.DistanceFunction distanceFunction) + { + PartitionKeyDefinition partitionKeyDefinition = CreateHashPartitionKey("/PartitionKey"); + + Cosmos.VectorEmbeddingPolicy vectorEmbeddingPolicy = new Cosmos.VectorEmbeddingPolicy(new Collection + { + new Cosmos.Embedding + { + Path = "/embedding", + DataType = Cosmos.VectorDataType.Float32, + Dimensions = 8, + DistanceFunction = distanceFunction + } + }); + + string queryText = @"SELECT TOP 10 c.title AS Title, VectorDistance(c.embedding, @vectorEmbedding, true) AS SimilarityScore + FROM c + ORDER BY VectorDistance(c.embedding, @vectorEmbedding, true)"; + + SqlQuerySpec sqlQuerySpec = new SqlQuerySpec( + queryText, + new SqlParameterCollection(new SqlParameter[] { new SqlParameter("@vectorEmbedding", VectorEmbedding) })); + + return new QueryPlanBaselineTestInput( + description, + partitionKeyDefinition, + vectorEmbeddingPolicy, + sqlQuerySpec, + Cosmos.GeospatialType.Geography); + } + + private static readonly double[] VectorEmbedding = new double[] { + 0.0039695268496870995, + 0.027338456362485886, + -0.005676387343555689, + -0.013547309674322605, + -0.002445543883368373, + 0.01579204574227333, + -0.016796082258224487, + -0.012471556663513184 }; private static PartitionKeyDefinition CreateHashPartitionKey( params string[] partitionKeys) => new PartitionKeyDefinition() @@ -1439,7 +1494,8 @@ public override QueryPlanBaselineTestOutput ExecuteTest(QueryPlanBaselineTestInp { TryCatch info = QueryPartitionProviderTestInstance.Object.TryGetPartitionedQueryExecutionInfoInternal( JsonConvert.SerializeObject(input.SqlQuerySpec), - input.PartitionKeyDefinition, + input.PartitionKeyDefinition, + input.VectorEmbeddingPolicy, requireFormattableOrderByQuery: true, isContinuationExpected: false, allowNonValueAggregateQuery: true, @@ -1459,18 +1515,36 @@ public override QueryPlanBaselineTestOutput ExecuteTest(QueryPlanBaselineTestInp public sealed class QueryPlanBaselineTestInput : BaselineTestInput { - internal PartitionKeyDefinition PartitionKeyDefinition { get; set; } - internal SqlQuerySpec SqlQuerySpec { get; set; } - internal Cosmos.GeospatialType? GeospatialType { get; set; } + internal PartitionKeyDefinition PartitionKeyDefinition { get; } + + internal Cosmos.VectorEmbeddingPolicy VectorEmbeddingPolicy { get; } + + internal SqlQuerySpec SqlQuerySpec { get; } + + internal Cosmos.GeospatialType? GeospatialType { get; } internal QueryPlanBaselineTestInput( string description, PartitionKeyDefinition partitionKeyDefinition, SqlQuerySpec sqlQuerySpec) - : base(description) + : this(description, partitionKeyDefinition, vectorEmbeddingPolicy: null, sqlQuerySpec, geospatialType: null) { this.PartitionKeyDefinition = partitionKeyDefinition; this.SqlQuerySpec = sqlQuerySpec; + } + + internal QueryPlanBaselineTestInput( + string description, + PartitionKeyDefinition partitionKeyDefinition, + Cosmos.VectorEmbeddingPolicy vectorEmbeddingPolicy, + SqlQuerySpec sqlQuerySpec, + Cosmos.GeospatialType? geospatialType) + : base(description) + { + this.PartitionKeyDefinition = partitionKeyDefinition; + this.VectorEmbeddingPolicy = vectorEmbeddingPolicy; + this.SqlQuerySpec = sqlQuerySpec; + this.GeospatialType = geospatialType; } public override void SerializeAsXml(XmlWriter xmlWriter) diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPlanRetrieverTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPlanRetrieverTests.cs index c1dc2fc3ae..7dd2de521d 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPlanRetrieverTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPlanRetrieverTests.cs @@ -31,7 +31,8 @@ public async Task ServiceInterop_BadRequestContainsInnerException() queryClient.Setup(c => c.TryGetPartitionedQueryExecutionInfoAsync( It.IsAny(), It.IsAny(), - It.IsAny(), + It.IsAny(), + It.IsAny(), It.IsAny(), It.IsAny(), It.IsAny(), @@ -45,7 +46,8 @@ public async Task ServiceInterop_BadRequestContainsInnerException() queryClient.Object, new SqlQuerySpec("selectttttt * from c"), ResourceType.Document, - new Documents.PartitionKeyDefinition() { Paths = new Collection() { "/id" } }, + new Documents.PartitionKeyDefinition() { Paths = new Collection() { "/id" } }, + vectorEmbeddingPolicy:null, hasLogicalPartitionKey: false, geospatialType: Cosmos.GeospatialType.Geography, useSystemPrefix: false, @@ -66,7 +68,8 @@ public async Task ServiceInterop_BadRequestContainsOriginalCosmosException() queryClient.Setup(c => c.TryGetPartitionedQueryExecutionInfoAsync( It.IsAny(), It.IsAny(), - It.IsAny(), + It.IsAny(), + It.IsAny(), It.IsAny(), It.IsAny(), It.IsAny(), @@ -81,7 +84,8 @@ public async Task ServiceInterop_BadRequestContainsOriginalCosmosException() queryClient.Object, new SqlQuerySpec("selectttttt * from c"), ResourceType.Document, - new Documents.PartitionKeyDefinition() { Paths = new Collection() { "/id" } }, + new Documents.PartitionKeyDefinition() { Paths = new Collection() { "/id" } }, + vectorEmbeddingPolicy: null, hasLogicalPartitionKey: false, geospatialType: Cosmos.GeospatialType.Geography, useSystemPrefix: false, @@ -100,7 +104,8 @@ public async Task ServiceInterop_E_UNEXPECTED() queryClient.Setup(c => c.TryGetPartitionedQueryExecutionInfoAsync( It.IsAny(), It.IsAny(), - It.IsAny(), + It.IsAny(), + It.IsAny(), It.IsAny(), It.IsAny(), It.IsAny(), @@ -115,7 +120,8 @@ public async Task ServiceInterop_E_UNEXPECTED() new SqlQuerySpec("Super secret query that triggers bug"), ResourceType.Document, new Documents.PartitionKeyDefinition() { Paths = new Collection() { "/id" } }, - hasLogicalPartitionKey: false, + vectorEmbeddingPolicy: null, + hasLogicalPartitionKey: false, geospatialType: Cosmos.GeospatialType.Geography, useSystemPrefix: false, NoOpTrace.Singleton)); diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/SubpartitionTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/SubpartitionTests.cs index 1605c09a1a..2308ffa424 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/SubpartitionTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/SubpartitionTests.cs @@ -171,7 +171,8 @@ internal static Tuple Get QueryPartitionProvider queryPartitionProvider = CreateCustomQueryPartitionProvider(); TryCatch tryGetQueryPlan = queryPartitionProvider.TryGetPartitionedQueryExecutionInfo( querySpecJsonString: querySpecJsonString, - partitionKeyDefinition: pkDefinition, + partitionKeyDefinition: pkDefinition, + vectorEmbeddingPolicy: null, requireFormattableOrderByQuery: true, isContinuationExpected: true, allowNonValueAggregateQuery: true, @@ -325,7 +326,8 @@ public override Task GetCachedContainerQueryProperties true, true) }, - SubpartitionTests.CreatePartitionKeyDefinition(), + SubpartitionTests.CreatePartitionKeyDefinition(), + vectorEmbeddingPolicy: null, Cosmos.GeospatialType.Geometry)); } @@ -349,14 +351,26 @@ public override Task> TryGetOverlappingRangesAs throw new NotImplementedException(); } - public override async Task> TryGetPartitionedQueryExecutionInfoAsync(SqlQuerySpec sqlQuerySpec, ResourceType resourceType, PartitionKeyDefinition partitionKeyDefinition, bool requireFormattableOrderByQuery, bool isContinuationExpected, bool allowNonValueAggregateQuery, bool hasLogicalPartitionKey, bool allowDCount, bool useSystemPrefix, Cosmos.GeospatialType geospatialType, CancellationToken cancellationToken) + public override Task> TryGetPartitionedQueryExecutionInfoAsync( + SqlQuerySpec sqlQuerySpec, + ResourceType resourceType, + PartitionKeyDefinition partitionKeyDefinition, + Cosmos.VectorEmbeddingPolicy vectorEmbeddingPolicy, + bool requireFormattableOrderByQuery, + bool isContinuationExpected, + bool allowNonValueAggregateQuery, + bool hasLogicalPartitionKey, + bool allowDCount, + bool useSystemPrefix, + Cosmos.GeospatialType geospatialType, + CancellationToken cancellationToken) { CosmosSerializerCore serializerCore = new(); using StreamReader streamReader = new(serializerCore.ToStreamSqlQuerySpec(sqlQuerySpec, Documents.ResourceType.Document)); string sqlQuerySpecJsonString = streamReader.ReadToEnd(); (PartitionedQueryExecutionInfo partitionedQueryExecutionInfo, QueryPartitionProvider queryPartitionProvider) = OptimisticDirectExecutionQueryBaselineTests.GetPartitionedQueryExecutionInfoAndPartitionProvider(sqlQuerySpecJsonString, partitionKeyDefinition); - return TryCatch.FromResult(partitionedQueryExecutionInfo); + return Task.FromResult(TryCatch.FromResult(partitionedQueryExecutionInfo)); } } } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Routing/PartitionRoutingHelperTest.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Routing/PartitionRoutingHelperTest.cs index 792621f54a..ca444d7a09 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Routing/PartitionRoutingHelperTest.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Routing/PartitionRoutingHelperTest.cs @@ -712,7 +712,8 @@ private static async Task PrefixPartitionKeyTestRunnerAsync( TryCatch tryGetQueryPlan = QueryPartitionProvider.TryGetPartitionedQueryExecutionInfo( querySpecJsonString: JsonConvert.SerializeObject(new SqlQuerySpec(queryText)), - partitionKeyDefinition: partitionKeyDefinition, + partitionKeyDefinition: partitionKeyDefinition, + vectorEmbeddingPolicy: null, requireFormattableOrderByQuery: true, isContinuationExpected: true, allowNonValueAggregateQuery: false, @@ -862,7 +863,8 @@ FROM c allowDCount: false, allowNonValueAggregates: true, useSystemPrefix: false, - partitionKeyDefinition: new PartitionKeyDefinition { Paths = new Collection { testcase.PartitionKey }, Kind = PartitionKind.Hash }, + partitionKeyDefinition: new PartitionKeyDefinition { Paths = new Collection { testcase.PartitionKey }, Kind = PartitionKind.Hash }, + vectorEmbeddingPolicy: null, queryPartitionProvider: QueryPartitionProviderTestInstance.Object, clientApiVersion: testcase.ClientApiVersion, geospatialType: Cosmos.GeospatialType.Geography, diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Tracing/TraceWriterBaselineTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Tracing/TraceWriterBaselineTests.cs index 39082e940d..c06d7c57dd 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Tracing/TraceWriterBaselineTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Tracing/TraceWriterBaselineTests.cs @@ -769,7 +769,8 @@ private static QueryInfo GetQueryPlan(string query) { TryCatch info = QueryPartitionProviderTestInstance.Object.TryGetPartitionedQueryExecutionInfoInternal( Newtonsoft.Json.JsonConvert.SerializeObject(new SqlQuerySpec(query)), - partitionKeyDefinition, + partitionKeyDefinition, + vectorEmbeddingPolicy: null, requireFormattableOrderByQuery: true, isContinuationExpected: false, allowNonValueAggregateQuery: true,