Add support for the optimized query plan that skips the order by rewrite

Author: neildsh

Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/HybridSearch/HybridSearchCrossPartitionQueryPipelineStage.cs

@@ -404,12 +404,10 @@ private static IReadOnlyList<ComponentWeight> ExtractComponentWeights(HybridSear
             for (int index = 0; index < hybridSearchQueryInfo.ComponentQueryInfos.Count; ++index)
             {
                 QueryInfo queryInfo = hybridSearchQueryInfo.ComponentQueryInfos[index];
-                Debug.Assert(queryInfo.HasOrderBy, "The component query should have an order by");
-                Debug.Assert(queryInfo.HasNonStreamingOrderBy, "The component query is a non streaming order by");
-                Debug.Assert(queryInfo.OrderBy.Count == 1, "The component query should have exactly one order by expression");
+                SortOrder sortOrder = queryInfo.HasOrderBy ? queryInfo.OrderBy[0] : SortOrder.Descending;
 
                 double componentWeight = useDefaultComponentWeight ? 1.0 : hybridSearchQueryInfo.ComponentWeights[index];
-                result.Add(new ComponentWeight(componentWeight, queryInfo.OrderBy[0]));
+                result.Add(new ComponentWeight(componentWeight, sortOrder));
             }
 
             return result;
@@ -635,14 +633,19 @@ private static void ComputeRrfScores(
 
         private static QueryInfo RewriteOrderByQueryInfo(QueryInfo queryInfo, GlobalFullTextSearchStatistics statistics, int componentCount)
         {
-            Debug.Assert(queryInfo.HasOrderBy, "The component query should have an order by");
-            Debug.Assert(queryInfo.HasNonStreamingOrderBy, "The component query is a non streaming order by");
+            IReadOnlyList<string> rewrittenOrderByExpressions = queryInfo.OrderByExpressions;
 
-            List<string> rewrittenOrderByExpressions = new List<string>(queryInfo.OrderByExpressions.Count);
-            foreach (string orderByExpression in queryInfo.OrderByExpressions)
+            if (queryInfo.HasOrderBy)
             {
-                string rewrittenOrderByExpression = FormatComponentQueryTextWorkaround(orderByExpression, statistics, componentCount);
-                rewrittenOrderByExpressions.Add(rewrittenOrderByExpression);
+                Debug.Assert(queryInfo.HasNonStreamingOrderBy, "The component query is a non streaming order by");
+                List<string> orderByExpressions = new List<string>(queryInfo.OrderByExpressions.Count);
+                foreach (string orderByExpression in queryInfo.OrderByExpressions)
+                {
+                    string rewrittenOrderByExpression = FormatComponentQueryTextWorkaround(orderByExpression, statistics, componentCount);
+                    orderByExpressions.Add(rewrittenOrderByExpression);
+                }
+
+                rewrittenOrderByExpressions = orderByExpressions;
             }
 
             string rewrittenQuery = FormatComponentQueryTextWorkaround(queryInfo.RewrittenQuery, statistics, componentCount);
@@ -777,18 +780,15 @@ private static string FormatComponentQueryTextWorkaround(string format, GlobalFu
 
         private class ComponentWeight
         {
-            public SortOrder SortOrder { get; }
-
             public double Weight { get; }
 
             public Comparison<double> Comparison { get; }
 
             public ComponentWeight(double weight, SortOrder sortOrder)
             {
                 this.Weight = weight;
-                this.SortOrder = sortOrder;
 
-                int comparisonFactor = (this.SortOrder == SortOrder.Ascending) ? 1 : -1;
+                int comparisonFactor = (sortOrder == SortOrder.Ascending) ? 1 : -1;
                 this.Comparison = (x, y) => comparisonFactor * x.CompareTo(y);
             }
         }

Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/HybridSearch/HybridSearchQueryResult.cs

@@ -54,17 +54,26 @@ public static HybridSearchQueryResult Create(CosmosElement document)
                 throw new ArgumentException($"{FieldNames.Payload} must exist.");
             }
 
-            if (!outerPayload.TryGetValue(FieldNames.Payload, out CosmosObject innerPayload))
+            if (outerPayload.TryGetValue(FieldNames.Payload, out CosmosObject innerPayload))
             {
-                throw new ArgumentException($"{FieldNames.Payload} must exist nested within the outer payload field.");
+                // Using the older format where the payload is nested.
+                if (!outerPayload.TryGetValue(FieldNames.ComponentScores, out CosmosArray componentScores))
+                {
+                    throw new ArgumentException($"{FieldNames.ComponentScores} must exist.");
+                }
+
+                return new HybridSearchQueryResult(rid, componentScores, innerPayload);
             }
 
-            if (!outerPayload.TryGetValue(FieldNames.ComponentScores, out CosmosArray componentScores))
             {
-                throw new ArgumentException($"{FieldNames.ComponentScores} must exist.");
-            }
+                // Using the newer format where the payload is not nested.
+                if (!cosmosObject.TryGetValue(FieldNames.ComponentScores, out CosmosArray componentScores))
+                {
+                    throw new ArgumentException($"{FieldNames.ComponentScores} must exist.");
+                }
 
-            return new HybridSearchQueryResult(rid, componentScores, innerPayload);
+                return new HybridSearchQueryResult(rid, componentScores, outerPayload);
+            }
         }
 
         private static class FieldNames

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

@@ -32,7 +32,8 @@ internal static class QueryPlanRetriever
             | QueryFeatures.NonStreamingOrderBy
             | QueryFeatures.CountIf
             | QueryFeatures.HybridSearch
-            | QueryFeatures.WeightedRankFusion;
+            | QueryFeatures.WeightedRankFusion
+            | QueryFeatures.HybridSearchSkipOrderByRewrite;
 
         private static readonly QueryFeatures SupportedQueryFeaturesWithoutNonStreamingOrderBy =
             SupportedQueryFeatures & (~QueryFeatures.NonStreamingOrderBy);