Resolving comments, adding new test cases

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/NonStreamingOrderByDocumentQueryExecutionContext.java

@@ -59,8 +59,6 @@ public class NonStreamingOrderByDocumentQueryExecutionContext
     private final Collection<ClientSideRequestStatistics> clientSideRequestStatistics;
     private Flux<OrderByRowResult<Document>> orderByObservable;
 
-    private int maxPageSizePerPartition;
-
     public NonStreamingOrderByDocumentQueryExecutionContext(
         DiagnosticsClientContext diagnosticsClientContext,
         IDocumentQueryClient client,
@@ -170,7 +168,7 @@ protected NonStreamingOrderByDocumentProducer createDocumentProducer(
             createRetryPolicyFunc,
             Document.class,
             correlatedActivityId,
-            maxPageSizePerPartition,
+            initialPageSize,
             continuationToken,
             top,
             this.getOperationContextTextProvider());

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/NonStreamingOrderByUtils.java

@@ -78,7 +78,6 @@ public Flux<OrderByRowResult<Document>> apply(Flux<DocumentProducer<Document>.Do
             // item and then remove the element. If we don't do this, then when adding this element the size of the pq
             // will be increased automatically by 50% and then there would be inconsistent results for later pages.
             PriorityBlockingQueue<OrderByRowResult<Document>> priorityQueue = new PriorityBlockingQueue<>(initialPageSize + 1, consumeComparer);
-
             return source.flatMap(documentProducerFeedResponse -> {
                     clientSideRequestStatistics.addAll(
                         diagnosticsAccessor.getClientSideRequestStatisticsForQueryPipelineAggregations(documentProducerFeedResponse
@@ -94,19 +93,19 @@ public Flux<OrderByRowResult<Document>> apply(Flux<DocumentProducer<Document>.Do
                             null);
                         priorityQueue.add(orderByRowResult);
                         if (priorityQueue.size() > initialPageSize) {
-                            PriorityBlockingQueue<OrderByRowResult<Document>> resultPriorityQueue = new PriorityBlockingQueue<>(initialPageSize + 1, consumeComparer);
-                            for (int i=0;i<initialPageSize;i++) {
-                                resultPriorityQueue.add(priorityQueue.poll());
+                            PriorityBlockingQueue<OrderByRowResult<Document>> tempPriorityQueue = new PriorityBlockingQueue<>(initialPageSize + 1, consumeComparer);
+                            for (int i = 0; i < initialPageSize; i++) {
+                                tempPriorityQueue.add(priorityQueue.poll());
                             }
                             priorityQueue.clear();
-                            priorityQueue.addAll(resultPriorityQueue);
+                            priorityQueue.addAll(tempPriorityQueue);
                         }
                     });
                     tracker.addCharge(documentProducerFeedResponse.pageResult.getRequestCharge());
                     // Returning an empty Flux since we are only processing and managing state here
                     return Flux.empty();
                 }, 1)
-                .thenMany(Flux.defer(()-> Flux.fromIterable(priorityQueue)));
+                .thenMany(Flux.defer(() -> Flux.fromIterable(priorityQueue)));
         }
     }
 }

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/QueryPlanRetriever.java

@@ -34,6 +34,11 @@ class QueryPlanRetriever {
         ImplementationBridgeHelpers.CosmosQueryRequestOptionsHelper.getCosmosQueryRequestOptionsAccessor();
 
     private static final String TRUE = "True";
+
+    // For a limited time, if the query runs against a region or emulator that has not yet been updated with the
+    // new NonStreamingOrderBy query feature the client might run into some issue of not being able to recognize this,
+    // and throw a 400 exception. If the environment variable `AZURE_COSMOS_DISABLE_NON_STREAMING_ORDER_BY` is set to
+    // True to opt out of this new query feature, we will return the OLD query features to operate correctly.
     private static final String SUPPORTED_QUERY_FEATURES = QueryFeature.Aggregate.name() + ", " +
                                                                QueryFeature.CompositeAggregate.name() + ", " +
                                                                QueryFeature.MultipleOrderBy.name() + ", " +

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosQueryRequestOptions.java

@@ -265,17 +265,6 @@ Integer getMaxItemCountForVectorSearch() {
         return this.actualRequestOptions.getMaxItemCountForVectorSearch();
     }
 
-    /**
-     * Sets the maximum item size to fetch during non-streaming order by queries.
-     *
-     * @param maxItemCountForVectorSearch the max number of items for vector search.
-     * @return the CosmosQueryRequestOptions.
-     */
-    CosmosQueryRequestOptions setMaxItemCountForVectorSearch(Integer maxItemCountForVectorSearch) {
-        this.actualRequestOptions.setMaxItemCountForVectorSearch(maxItemCountForVectorSearch);
-        return this;
-    }
-
     /**
      * Gets the request continuation token.
      *

sdk/cosmos/test-resources.json

@@ -117,7 +117,11 @@
           "maxStalenessPrefix": 100
         },
         "locations": "[variables('locationsConfiguration')]",
-        "capabilities": [],
+        "capabilities": [
+          {
+            "name": "EnableNoSQLVectorSearch"
+          }
+        ],
         "ipRules": []
       }
     }