Query: Fixes partition range evaluation for spatial queries (#3495)

* Initial commit

* Update.

* Pull/Rebase

* Addressed comments.

* Build fix

Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs

@@ -248,6 +248,7 @@ private static async Task<TryCatch<IQueryPipelineStage>> TryCreateCoreContextAsy
                             cosmosQueryContext.ResourceTypeEnum,
                             partitionKeyDefinition,
                             inputParameters.PartitionKey != null,
+                            containerQueryProperties.GeospatialType,
                             cosmosQueryContext.UseSystemPrefix,
                             createQueryPipelineTrace,
                             cancellationToken);

Microsoft.Azure.Cosmos/src/Query/Core/QueryClient/ContainerQueryProperties.cs

@@ -11,15 +11,18 @@ internal readonly struct ContainerQueryProperties
         public ContainerQueryProperties(
             string resourceId,
             string effectivePartitionKeyString,
-            PartitionKeyDefinition partitionKeyDefinition)
+            PartitionKeyDefinition partitionKeyDefinition,
+            Cosmos.GeospatialType geospatialType)
         {
             this.ResourceId = resourceId;
             this.EffectivePartitionKeyString = effectivePartitionKeyString;
             this.PartitionKeyDefinition = partitionKeyDefinition;
+            this.GeospatialType = geospatialType;
         }
 
         public string ResourceId { get; }
         public string EffectivePartitionKeyString { get; }
         public PartitionKeyDefinition PartitionKeyDefinition { get; }
+        public Cosmos.GeospatialType GeospatialType { get; }
     }
 }

Microsoft.Azure.Cosmos/src/Query/Core/QueryClient/CosmosQueryClient.cs

@@ -47,6 +47,7 @@ public abstract Task<TryCatch<PartitionedQueryExecutionInfo>> TryGetPartitionedQ
             bool hasLogicalPartitionKey,
             bool allowDCount,
             bool useSystemPrefix,
+            Cosmos.GeospatialType geospatialType,
             CancellationToken cancellationToken);
 
         public abstract Task<TryCatch<QueryPage>> ExecuteItemQueryAsync(

Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryPartitionProvider.cs

@@ -118,7 +118,8 @@ public TryCatch<PartitionedQueryExecutionInfo> TryGetPartitionedQueryExecutionIn
             bool allowNonValueAggregateQuery,
             bool hasLogicalPartitionKey,
             bool allowDCount,
-            bool useSystemPrefix)
+            bool useSystemPrefix,
+            GeospatialType geospatialType)
         {
             TryCatch<PartitionedQueryExecutionInfoInternal> tryGetInternalQueryInfo = this.TryGetPartitionedQueryExecutionInfoInternal(
                 querySpecJsonString: querySpecJsonString,
@@ -128,7 +129,8 @@ public TryCatch<PartitionedQueryExecutionInfo> TryGetPartitionedQueryExecutionIn
                 allowNonValueAggregateQuery: allowNonValueAggregateQuery,
                 hasLogicalPartitionKey: hasLogicalPartitionKey,
                 allowDCount: allowDCount,
-                useSystemPrefix: useSystemPrefix);
+                useSystemPrefix: useSystemPrefix,
+                geospatialType: geospatialType);
             if (!tryGetInternalQueryInfo.Succeeded)
             {
                 return TryCatch<PartitionedQueryExecutionInfo>.FromException(tryGetInternalQueryInfo.Exception);
@@ -169,7 +171,8 @@ internal TryCatch<PartitionedQueryExecutionInfoInternal> TryGetPartitionedQueryE
             bool allowNonValueAggregateQuery,
             bool hasLogicalPartitionKey,
             bool allowDCount,
-            bool useSystemPrefix)
+            bool useSystemPrefix,
+            GeospatialType geospatialType)
         {
             if (querySpecJsonString == null || partitionKeyDefinition == null)
             {
@@ -200,7 +203,6 @@ internal TryCatch<PartitionedQueryExecutionInfoInternal> TryGetPartitionedQueryE
             }
 
             PartitionKind partitionKind = partitionKeyDefinition.Kind;
-            GeospatialType defaultGeopatialType = GeospatialType.Geography;
 
             this.Initialize();
 
@@ -219,7 +221,7 @@ internal TryCatch<PartitionedQueryExecutionInfoInternal> TryGetPartitionedQueryE
                         bIsContinuationExpected = Convert.ToInt32(isContinuationExpected),
                         bRequireFormattableOrderByQuery = Convert.ToInt32(requireFormattableOrderByQuery),
                         bUseSystemPrefix = Convert.ToInt32(useSystemPrefix),
-                        eGeospatialType = Convert.ToInt32(defaultGeopatialType),
+                        eGeospatialType = Convert.ToInt32(geospatialType),
                         ePartitionKind = Convert.ToInt32(partitionKind)
                     };
 

Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryPlanHandler.cs

@@ -29,6 +29,7 @@ public async Task<TryCatch<PartitionedQueryExecutionInfo>> TryGetQueryPlanAsync(
             QueryFeatures supportedQueryFeatures,
             bool hasLogicalPartitionKey,
             bool useSystemPrefix,
+            GeospatialType geospatialType,
             CancellationToken cancellationToken)
         {
             cancellationToken.ThrowIfCancellationRequested();
@@ -49,6 +50,7 @@ public async Task<TryCatch<PartitionedQueryExecutionInfo>> TryGetQueryPlanAsync(
                 partitionKeyDefinition,
                 hasLogicalPartitionKey,
                 useSystemPrefix,
+                geospatialType,
                 cancellationToken);
             if (!tryGetQueryInfo.Succeeded)
             {
@@ -76,6 +78,7 @@ public async Task<TryCatch<PartitionedQueryExecutionInfo>> TryGetQueryPlanAsync(
             PartitionKeyDefinition partitionKeyDefinition,
             bool hasLogicalPartitionKey,
             bool useSystemPrefix,
+            GeospatialType geospatialType,
             CancellationToken cancellationToken = default)
         {
             if (sqlQuerySpec == null)
@@ -96,6 +99,7 @@ public async Task<TryCatch<PartitionedQueryExecutionInfo>> TryGetQueryPlanAsync(
                 partitionKeyDefinition,
                 hasLogicalPartitionKey,
                 useSystemPrefix,
+                geospatialType,
                 cancellationToken);
             if (tryGetQueryInfo.Failed)
             {
@@ -114,6 +118,7 @@ private Task<TryCatch<PartitionedQueryExecutionInfo>> TryGetQueryInfoAsync(
             PartitionKeyDefinition partitionKeyDefinition,
             bool hasLogicalPartitionKey,
             bool useSystemPrefix,
+            Cosmos.GeospatialType geospatialType,
             CancellationToken cancellationToken = default)
         {
             cancellationToken.ThrowIfCancellationRequested();
@@ -128,6 +133,7 @@ private Task<TryCatch<PartitionedQueryExecutionInfo>> TryGetQueryInfoAsync(
                 hasLogicalPartitionKey: hasLogicalPartitionKey,
                 allowDCount: true,
                 useSystemPrefix: useSystemPrefix,
+                geospatialType: geospatialType,
                 cancellationToken: cancellationToken);
         }
 

Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryPlanRetriever.cs

@@ -38,6 +38,7 @@ public static async Task<PartitionedQueryExecutionInfo> GetQueryPlanWithServiceI
             Documents.ResourceType resourceType,
             PartitionKeyDefinition partitionKeyDefinition,
             bool hasLogicalPartitionKey,
+            GeospatialType geospatialType,
             bool useSystemPrefix,
             ITrace trace,
             CancellationToken cancellationToken = default)
@@ -70,6 +71,7 @@ public static async Task<PartitionedQueryExecutionInfo> GetQueryPlanWithServiceI
                     QueryPlanRetriever.SupportedQueryFeatures,
                     hasLogicalPartitionKey,
                     useSystemPrefix,
+                    geospatialType,
                     cancellationToken);
 
                 if (!tryGetQueryPlan.Succeeded)

Microsoft.Azure.Cosmos/src/Query/v2Query/DefaultDocumentQueryExecutionContext.cs

@@ -319,6 +319,7 @@ private static bool ServiceInteropAvailable()
                         partitionKeyDefinition: partitionKeyDefinition,
                         queryPartitionProvider: queryPartitionProvider,
                         clientApiVersion: version,
+                        geospatialType: collection.GeospatialConfig.GeospatialType,
                         out QueryInfo _);
                 }
                 else if (request.Properties != null && request.Properties.TryGetValue(

Microsoft.Azure.Cosmos/src/Query/v2Query/DocumentQueryExecutionContextBase.cs

@@ -160,6 +160,7 @@ public async Task<PartitionedQueryExecutionInfo> GetPartitionedQueryExecutionInf
             bool allowNonValueAggregateQuery,
             bool hasLogicalPartitionKey,
             bool allowDCount,
+            Cosmos.GeospatialType geospatialType,
             CancellationToken cancellationToken)
         {
             cancellationToken.ThrowIfCancellationRequested();
@@ -174,6 +175,7 @@ public async Task<PartitionedQueryExecutionInfo> GetPartitionedQueryExecutionInf
                 allowNonValueAggregateQuery: allowNonValueAggregateQuery,
                 hasLogicalPartitionKey: hasLogicalPartitionKey,
                 allowDCount: allowDCount,
+                geospatialType: geospatialType,
                 useSystemPrefix: false);
             if (!tryGetPartitionedQueryExecutionInfo.Succeeded)
             {

Microsoft.Azure.Cosmos/src/Query/v2Query/DocumentQueryExecutionContextFactory.cs

@@ -104,6 +104,7 @@ public static async Task<IDocumentQueryExecutionContext> CreateDocumentQueryExec
                     allowNonValueAggregateQuery: true,
                     hasLogicalPartitionKey: feedOptions.PartitionKey != null,
                     allowDCount: true,
+                    geospatialType: collection.GeospatialConfig.GeospatialType,
                     cancellationToken: token);
 
                 if (DocumentQueryExecutionContextFactory.ShouldCreateSpecializedDocumentQueryExecutionContext(

Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs

@@ -74,7 +74,8 @@ public override async Task<ContainerQueryProperties> GetCachedContainerQueryProp
             return new ContainerQueryProperties(
                 containerProperties.ResourceId,
                 effectivePartitionKeyString,
-                containerProperties.PartitionKey);
+                containerProperties.PartitionKey,
+                containerProperties.GeospatialConfig.GeospatialType);
         }
 
         public override async Task<TryCatch<PartitionedQueryExecutionInfo>> TryGetPartitionedQueryExecutionInfoAsync(
@@ -87,6 +88,7 @@ public override async Task<TryCatch<PartitionedQueryExecutionInfo>> TryGetPartit
             bool hasLogicalPartitionKey,
             bool allowDCount,
             bool useSystemPrefix,
+            Cosmos.GeospatialType geospatialType,
             CancellationToken cancellationToken)
         {
             string queryString = null;
@@ -109,7 +111,8 @@ public override async Task<TryCatch<PartitionedQueryExecutionInfo>> TryGetPartit
                 allowNonValueAggregateQuery: allowNonValueAggregateQuery,
                 hasLogicalPartitionKey: hasLogicalPartitionKey,
                 allowDCount: allowDCount,
-                useSystemPrefix: useSystemPrefix);
+                useSystemPrefix: useSystemPrefix,
+                geospatialType: geospatialType);
         }
 
         public override async Task<TryCatch<QueryPage>> ExecuteItemQueryAsync(

Microsoft.Azure.Cosmos/src/Resource/Container/ContainerCore.Items.cs

@@ -351,6 +351,7 @@ public override async Task<TryExecuteQueryResult> TryExecuteQueryAsync(
             string continuationToken,
             FeedRangeInternal feedRangeInternal,
             QueryRequestOptions requestOptions,
+            GeospatialType geospatialType,
             CancellationToken cancellationToken = default)
         {
             if (queryDefinition == null)
@@ -418,6 +419,7 @@ public override async Task<TryExecuteQueryResult> TryExecuteQueryAsync(
                 partitionKeyDefinition,
                 requestOptions.PartitionKey.HasValue,
                 useSystemPrefix: QueryIterator.IsSystemPrefixExpected(requestOptions),
+                geospatialType: geospatialType,
                 cancellationToken);
 
             if (tryGetQueryInfoAndIfSupported.Failed)

Microsoft.Azure.Cosmos/src/Resource/Container/ContainerInternal.cs

@@ -71,6 +71,7 @@ public abstract Task<TryExecuteQueryResult> TryExecuteQueryAsync(
             string continuationToken,
             FeedRangeInternal feedRangeInternal,
             QueryRequestOptions requestOptions,
+            GeospatialType geospatialType,
             CancellationToken cancellationToken = default);
 
         public abstract FeedIterator GetStandByFeedIterator(

Microsoft.Azure.Cosmos/src/Routing/PartitionRoutingHelper.cs

@@ -37,6 +37,7 @@ public static IReadOnlyList<Range<string>> GetProvidedPartitionKeyRanges(
             PartitionKeyDefinition partitionKeyDefinition,
             QueryPartitionProvider queryPartitionProvider,
             string clientApiVersion,
+            Cosmos.GeospatialType geospatialType,
             out QueryInfo queryInfo)
         {
             if (querySpecJsonString == null)
@@ -62,7 +63,8 @@ public static IReadOnlyList<Range<string>> GetProvidedPartitionKeyRanges(
                 allowNonValueAggregateQuery: allowNonValueAggregates,
                 hasLogicalPartitionKey: hasLogicalPartitionKey,
                 allowDCount: allowDCount,
-                useSystemPrefix: useSystemPrefix);
+                useSystemPrefix: useSystemPrefix,
+                geospatialType: geospatialType);
             if (!tryGetPartitionQueryExecutionInfo.Succeeded)
             {
                 throw new BadRequestException(tryGetPartitionQueryExecutionInfo.Exception);

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemTests.cs

@@ -1589,7 +1589,8 @@ public async Task ItemEpkQuerySingleKeyRangeValidation()
                 ContainerQueryProperties containerQueryProperties = new ContainerQueryProperties(
                     containerResponse.Resource.ResourceId,
                     null,
-                    containerResponse.Resource.PartitionKey);
+                    containerResponse.Resource.PartitionKey,
+                    containerResponse.Resource.GeospatialConfig.GeospatialType);
 
                 // There should only be one range since the EPK option is set.
                 List<PartitionKeyRange> partitionKeyRanges = await CosmosQueryExecutionContextFactory.GetTargetPartitionKeyRangesAsync(

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/FeedToken/QueryFeedTokenTests.cs

@@ -55,7 +55,8 @@ public async Task GetTargetPartitionKeyRangesAsyncWithFeedRange()
                 ContainerQueryProperties containerQueryProperties = new ContainerQueryProperties(
                     containerResponse.Resource.ResourceId,
                     null,
-                    containerResponse.Resource.PartitionKey);
+                    containerResponse.Resource.PartitionKey,
+                    containerResponse.Resource.GeospatialConfig.GeospatialType);
 
                 IReadOnlyList<FeedRange> feedTokens = await container.GetFeedRangesAsync();
 

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/PartitioningQueryTests.cs

@@ -489,5 +489,90 @@ async Task ImplementationAsync(
                 Assert.IsTrue(actualPartitionKeyValues.SetEquals(args.ExpectedPartitionKeyValues));
             }
         }
+
+        [TestMethod]
+        public async Task TestGeospatial()
+        {
+            string[] inputDocs = new[]
+            {
+                @"{""id"":""documentId1"",""key"":""A"",""prop"":3,""shortArray"":[{""a"":5}]}",
+                @"{""id"":""documentId2"",""key"":""A"",""prop"":2,""shortArray"":[{""a"":6}]}",
+                @"{""id"":""documentId3"",""key"":""A"",""prop"":1,""shortArray"":[{""a"":7}]}",
+                @"{""id"":""documentId4"",""key"":5,""prop"":3,""shortArray"":[{""a"":5}]}",
+                @"{""id"":""documentId5"",""key"":5,""prop"":2,""shortArray"":[{""a"":6}]}",
+                @"{""id"":""documentId6"",""key"":5,""prop"":1,""shortArray"":[{""a"":7}]}",
+                @"{""id"":""documentId10"",""prop"":3,""shortArray"":[{""a"":5}]}",
+                @"{""id"":""documentId11"",""prop"":2,""shortArray"":[{""a"":6}]}",
+                @"{""id"":""documentId12"",""prop"":1,""shortArray"":[{""a"":7}]}",
+            };
+
+            string All = "documentId4,documentId5,documentId6,documentId10,documentId11,documentId12,documentId1,documentId2,documentId3";
+            string None = string.Empty;
+            var testVariations = new[]
+            {
+                new
+                {
+                    Query = "SELECT c.id FROM c WHERE ST_DISTANCE({'type': 'Polygon', 'coordinates': [[[35, 10], [45, 45], [15, 40], [10, 20], [35, 10]], [[20, 30], [35, 35], [30, 20], [20, 30]]]}, {'type': 'Point', 'coordinates': [30, 10]}) > 66408.034483",
+                    Expected = new Dictionary<Cosmos.GeospatialType, string>
+                        {
+                            { Cosmos.GeospatialType.Geography, All },
+                            { Cosmos.GeospatialType.Geometry, None },
+                        }
+                },
+                new
+                {
+                    Query = "SELECT c.id FROM c WHERE ST_ISVALID({'type': 'Polygon', 'coordinates': [[[-1000, 1000], [1000, 1000], [1000, 4000], [-1000, 4000], [-1000, 1000]]]})",
+                    Expected = new Dictionary<Cosmos.GeospatialType, string>
+                        {
+                            { Cosmos.GeospatialType.Geography, None },
+                            { Cosmos.GeospatialType.Geometry, All },
+                        }
+                },
+                new
+                {
+                    Query = "SELECT * FROM c WHERE NOT ST_WITHIN({'type': 'Point', 'coordinates': [0, 40]}, {'type':'Polygon','coordinates':[[[-60,20], [70,20], [70,70], [-60,70], [-60,20]]]})",
+                    Expected = new Dictionary<Cosmos.GeospatialType, string>
+                        {
+                            { Cosmos.GeospatialType.Geography, All },
+                            { Cosmos.GeospatialType.Geometry, None },
+                        }
+                },
+            };
+
+            foreach (Cosmos.GeospatialType geospatialType in new[] { Cosmos.GeospatialType.Geography, Cosmos.GeospatialType.Geometry })
+            {
+                await this.CreateIngestQueryDeleteAsync(
+                    ConnectionModes.Direct,
+                    CollectionTypes.MultiPartition,
+                    inputDocs,
+                    ImplementationAsync,
+                    "/key",
+                    geospatialType: geospatialType);
+
+                async Task ImplementationAsync(Container container, IReadOnlyList<CosmosObject> documents)
+                {
+                    foreach (var testVariation in testVariations)
+                    {
+                        string expectedResult = string.Join(",", testVariation.Expected[geospatialType]);
+
+                        FeedIterator<Document> resultSetIterator = container.GetItemQueryIterator<Document>(
+                            queryText: testVariation.Query,
+                            requestOptions: new QueryRequestOptions());
+
+                        List<Document> result = new List<Document>();
+                        while (resultSetIterator.HasMoreResults)
+                        {
+                            result.AddRange(await resultSetIterator.ReadNextAsync());
+                        }
+
+                        string resultDocIds = string.Join(",", result.Select(doc => doc.Id));
+                        Assert.AreEqual(
+                            expectedResult,
+                            resultDocIds,
+                            $"{Environment.NewLine}Query failed for geospatial type '{geospatialType}'{Environment.NewLine}{testVariation.Query}");
+                    }
+                }
+            }
+        }
     }
 }