Initial changes

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/CosmosQueryRequestOptionsImpl.java

@@ -204,7 +204,7 @@ public CosmosQueryRequestOptionsImpl setMaxItemCount(Integer maxItemCount) {
      * @return the max number of items for vector search.
      */
     public Integer getMaxItemSizeForVectorSearch() {
-        return maxItemSizeForVectorSearch;
+        return this.maxItemSizeForVectorSearch;
     }
 
     /**

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

@@ -29,6 +29,7 @@
 
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.UUID;
@@ -125,10 +126,13 @@ private void initialize(
         // Since the continuation token will always be null,
         // we don't need to handle any initialization based on continuationToken.
         // We can directly initialize without any consideration for continuationToken.
+        Map<FeedRangeEpkImpl, String> partitionKeyRangeToContinuationToken = new HashMap<>();
+        for (FeedRangeEpkImpl feedRangeEpk : feedRanges) {
+            partitionKeyRangeToContinuationToken.put(feedRangeEpk,
+                null);
+        }
         super.initialize(collection,
-            feedRanges.stream().collect(Collectors.toMap(
-                feedRangeEpk -> feedRangeEpk,
-                feedRangeEpk -> null)),
+            partitionKeyRangeToContinuationToken,
             initialPageSize,
             new SqlQuerySpec(querySpec.getQueryText().replace(FormatPlaceHolder, True),
                 querySpec.getParameters()));

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

@@ -74,7 +74,10 @@ private PageToItemTransformer(RequestChargeTracker tracker, Map<String, QueryMet
 
         @Override
         public Flux<OrderByRowResult<Document>> apply(Flux<DocumentProducer<Document>.DocumentProducerFeedResponse> source) {
-            PriorityBlockingQueue<OrderByRowResult<Document>> priorityQueue = new PriorityBlockingQueue<>(initialPageSize, consumeComparer);
+            // the size of the priority queue is set to size+1, because when the pq reaches the max size we add that
+            // 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(

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

@@ -5,6 +5,7 @@
 import com.azure.cosmos.CosmosItemSerializer;
 import com.azure.cosmos.implementation.DiagnosticsClientContext;
 import com.azure.cosmos.implementation.DocumentCollection;
+import com.azure.cosmos.implementation.HttpConstants;
 import com.azure.cosmos.implementation.ImplementationBridgeHelpers;
 import com.azure.cosmos.implementation.Document;
 import com.azure.cosmos.implementation.ObjectNodeMap;
@@ -56,7 +57,11 @@ private static BiFunction<String, PipelinedDocumentQueryParams<Document>, Flux<I
                 CosmosQueryRequestOptions orderByCosmosQueryRequestOptions =
                     qryOptAccessor.clone(requestOptions);
                 if (queryInfo.hasNonStreamingOrderBy()) {
-                    checkNotNull(continuationToken, "Can not use a continuation token for a vector search query ");
+                    if (continuationToken != null) {
+                        throw new NonStreamingOrderByBadRequestException(
+                            HttpConstants.StatusCodes.BADREQUEST,
+                            "Can not use a continuation token for a vector search query");
+                    }
                     qryOptAccessor.getImpl(orderByCosmosQueryRequestOptions).setCustomItemSerializer(null);
                     documentQueryParams.setCosmosQueryRequestOptions(orderByCosmosQueryRequestOptions);
                     return NonStreamingOrderByDocumentQueryExecutionContext.createAsync(diagnosticsClientContext, client, documentQueryParams, collection);