Resolving comments

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -4,7 +4,7 @@
 
 #### Features Added
 
-* Added a new non-streaming OrderBy query pipeline and a query feature`NonStreamingOrderBy` to support Vector Search queries. - See [PR 39897](https://github.com/Azure/azure-sdk-for-java/pull/39897/) 
+* Added support for non-streaming OrderBy query and a query feature `NonStreamingOrderBy` to support Vector Search queries. - See [PR 39897](https://github.com/Azure/azure-sdk-for-java/pull/39897/) 
 
 #### Breaking Changes
 

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

@@ -166,6 +166,13 @@ public class Configs {
     public static final String MAX_RETRIES_IN_LOCAL_REGION_WHEN_REMOTE_REGION_PREFERRED = "COSMOS.MAX_RETRIES_IN_LOCAL_REGION_WHEN_REMOTE_REGION_PREFERRED";
     private static final int DEFAULT_MAX_RETRIES_IN_LOCAL_REGION_WHEN_REMOTE_REGION_PREFERRED = 1;
 
+    private static final String MAX_ITEM_SIZE_FOR_VECTOR_SEARCH = "COSMOS.MAX_ITEM_SIZE_FOR_VECTOR_SEARCH";
+    private static final int DEFAULT_MAX_ITEM_SIZE_FOR_VECTOR_SEARCH = 50000;
+
+    private static final String MAX_ITEM_SIZE_FOR_VECTOR_SEARCH_ENABLED = "COSMOS.MAX_ITEM_SIZE_FOR_VECTOR_SEARCH_ENABLED";
+
+    private static final boolean DEFAULT_MAX_ITEM_SIZE_FOR_VECTOR_SEARCH_ENABLED = true;
+
     public static final int MIN_MAX_RETRIES_IN_LOCAL_REGION_WHEN_REMOTE_REGION_PREFERRED = 1;
 
     public static final String TCP_CONNECTION_ACQUISITION_TIMEOUT_IN_MS = "COSMOS.TCP_CONNECTION_ACQUISITION_TIMEOUT_IN_MS";
@@ -484,6 +491,14 @@ public static int getMaxRetriesInLocalRegionWhenRemoteRegionPreferred() {
                 MIN_MAX_RETRIES_IN_LOCAL_REGION_WHEN_REMOTE_REGION_PREFERRED);
     }
 
+    public static int getMaxItemSizeForVectorSearch() {
+        return getJVMConfigAsInt(MAX_ITEM_SIZE_FOR_VECTOR_SEARCH, DEFAULT_MAX_ITEM_SIZE_FOR_VECTOR_SEARCH);
+    }
+
+    public static boolean getMaxItemSizeForVectorSearchEnabled() {
+        return getJVMConfigAsBoolean(MAX_ITEM_SIZE_FOR_VECTOR_SEARCH_ENABLED, DEFAULT_MAX_ITEM_SIZE_FOR_VECTOR_SEARCH_ENABLED);
+    }
+
     public static Duration getMinRetryTimeInLocalRegionWhenRemoteRegionPreferred() {
         return
             Duration.ofMillis(Math.max(

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

@@ -25,6 +25,7 @@ public final class CosmosQueryRequestOptionsImpl extends CosmosQueryRequestOptio
     private boolean queryPlanRetrievalDisallowed;
     private boolean emptyPageDiagnosticsEnabled;
     private String queryName;
+    private Integer maxItemSizeForVectorSearch;
     private List<CosmosDiagnostics> cancelledRequestDiagnosticsTracker = new ArrayList<>();
 
     /**
@@ -62,6 +63,7 @@ public CosmosQueryRequestOptionsImpl(CosmosQueryRequestOptionsImpl options) {
         this.queryName = options.queryName;
         this.feedRange = options.feedRange;
         this.cancelledRequestDiagnosticsTracker = options.cancelledRequestDiagnosticsTracker;
+        this.maxItemSizeForVectorSearch = options.maxItemSizeForVectorSearch;
     }
 
     /**
@@ -196,6 +198,26 @@ public CosmosQueryRequestOptionsImpl setMaxItemCount(Integer maxItemCount) {
         return this;
     }
 
+    /**
+     * Gets the maximum item size to fetch during non-streaming order by queries.
+     *
+     * @return the max number of items for vector search.
+     */
+    public Integer getMaxItemSizeForVectorSearch() {
+        return maxItemSizeForVectorSearch;
+    }
+
+    /**
+     * Sets the maximum item size to fetch during non-streaming order by queries.
+     *
+     * @param maxItemSizeForVectorSearch the max number of items for vector search.
+     * return the CosmosQueryRequestOptions.
+     */
+    public CosmosQueryRequestOptionsImpl setMaxItemSizeForVectorSearch(Integer maxItemSizeForVectorSearch) {
+        this.maxItemSizeForVectorSearch = maxItemSizeForVectorSearch;
+        return this;
+    }
+
     /**
      * Gets the request continuation token.
      *

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

@@ -4,6 +4,7 @@
 
 import com.azure.cosmos.BridgeInternal;
 import com.azure.cosmos.implementation.BadRequestException;
+import com.azure.cosmos.implementation.Configs;
 import com.azure.cosmos.implementation.Constants;
 import com.azure.cosmos.implementation.DiagnosticsClientContext;
 import com.azure.cosmos.implementation.DocumentCollection;
@@ -240,7 +241,8 @@ private static boolean canCacheQuery(QueryInfo queryInfo) {
                    && !queryInfo.hasTop()
                    && !queryInfo.hasOffset()
                    && !queryInfo.hasDCount()
-                   && !queryInfo.hasOrderBy();
+                   && !queryInfo.hasOrderBy()
+                   && !queryInfo.hasNonStreamingOrderBy();
     }
 
     private static boolean isScopedToSinglePartition(CosmosQueryRequestOptions cosmosQueryRequestOptions) {
@@ -360,19 +362,33 @@ public static <T> Flux<? extends IDocumentQueryExecutionContext<T>> createSpecia
         boolean getLazyFeedResponse = queryInfo.hasTop();
 
         // We need to compute the optimal initial age size for non-streaming order-by queries
-        if (queryInfo.hasNonStreamingOrderBy()) {
+        if (queryInfo.hasNonStreamingOrderBy() && Configs.getMaxItemSizeForVectorSearchEnabled()) {
             // Validate the TOP or LIMIT for non-streaming order-by queries
-            if (!queryInfo.hasTop() || !queryInfo.hasLimit() || queryInfo.getTop() < 0 || queryInfo.getLimit() < 0) {
+            if (!queryInfo.hasTop() && !queryInfo.hasLimit() && queryInfo.getTop() < 0 && queryInfo.getLimit() < 0) {
                 throw new NonStreamingOrderByBadRequestException(HttpConstants.StatusCodes.BADREQUEST,
                     "Executing a vector search query without TOP or LIMIT can consume a large number of RUs" +
                         "very fast and have long runtimes. Please ensure you are using one of the above two filters" +
                         "with you vector search query.");
             }
-
+            // Validate the size of TOP or LIMIT against MaxItemSizeForVectorSearch
+            int maxLimit = Math.max(queryInfo.hasTop() ? queryInfo.getTop() : 0,
+                queryInfo.hasLimit() ? queryInfo.getLimit() : 0);
+            int maxItemSizeForVectorSearch = Math.max(Configs.getMaxItemSizeForVectorSearch(),
+                ModelBridgeInternal.getMaxItemSizeForVectorSearchFromQueryRequestOptions(cosmosQueryRequestOptions));
+            if (maxLimit > maxItemSizeForVectorSearch) {
+                throw new NonStreamingOrderByBadRequestException(HttpConstants.StatusCodes.BADREQUEST,
+                    "Executing a vector search query with TOP or LIMIT larger than the maxItemSizeForVectorSearch " +
+                        "is not allowed");
+            }
             // Set initialPageSize based on the smallest of TOP or LIMIT
-            if (queryInfo.hasTop() || queryInfo.hasLimit() ) {
-                initialPageSize = Math.min(queryInfo.hasTop() ? queryInfo.getTop() : Integer.MAX_VALUE,
+            if (queryInfo.hasTop() || queryInfo.hasLimit()) {
+                int pageSizeWithTopOrLimit = Math.min(queryInfo.hasTop() ? queryInfo.getTop() : Integer.MAX_VALUE,
                                            queryInfo.hasLimit() ? queryInfo.getLimit() : Integer.MAX_VALUE);
+                if (initialPageSize > 0) {
+                    initialPageSize = Math.min(pageSizeWithTopOrLimit, initialPageSize);
+                } else {
+                    initialPageSize = pageSizeWithTopOrLimit;
+                }
             }
         }
 

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

@@ -71,23 +71,20 @@ public NonStreamingOrderByDocumentQueryExecutionContext(
         OrderbyRowComparer<Document> consumeComparer,
         UUID correlatedActivityId,
         boolean hasSelectValue,
-        final AtomicBoolean isQueryCancelledOnTimeout,
-        int maxPageSizePerPartition) {
+        final AtomicBoolean isQueryCancelledOnTimeout) {
         super(diagnosticsClientContext, client, resourceTypeEnum, Document.class, query, cosmosQueryRequestOptions,
             resourceLink, rewrittenQuery, correlatedActivityId, hasSelectValue, isQueryCancelledOnTimeout);
         this.consumeComparer = consumeComparer;
         this.tracker = new RequestChargeTracker();
         this.queryMetricMap = new ConcurrentHashMap<>();
         this.clientSideRequestStatistics = ConcurrentHashMap.newKeySet();
-        this.maxPageSizePerPartition = maxPageSizePerPartition;
     }
 
     public static Flux<IDocumentQueryExecutionComponent<Document>> createAsync(
         DiagnosticsClientContext diagnosticsClientContext,
         IDocumentQueryClient client,
         PipelinedDocumentQueryParams<Document> initParams,
-        DocumentCollection collection,
-        int maxPageSizePerPartition) {
+        DocumentCollection collection) {
 
         QueryInfo queryInfo = initParams.getQueryInfo();
 
@@ -102,8 +99,7 @@ public static Flux<IDocumentQueryExecutionComponent<Document>> createAsync(
             new OrderbyRowComparer<>(queryInfo.getOrderBy()),
             initParams.getCorrelatedActivityId(),
             queryInfo.hasSelectValue(),
-            initParams.isQueryCancelledOnTimeout(),
-            maxPageSizePerPartition);
+            initParams.isQueryCancelledOnTimeout());
 
         context.setTop(initParams.getTop());
 

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

@@ -1,22 +1,23 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
 package com.azure.cosmos.implementation.query;
 
 import com.azure.cosmos.BridgeInternal;
 import com.azure.cosmos.implementation.ClientSideRequestStatistics;
+import com.azure.cosmos.implementation.Configs;
 import com.azure.cosmos.implementation.Document;
 import com.azure.cosmos.implementation.ImplementationBridgeHelpers;
 import com.azure.cosmos.implementation.QueryMetrics;
 import com.azure.cosmos.implementation.RequestChargeTracker;
 import com.azure.cosmos.implementation.Resource;
 import com.azure.cosmos.implementation.query.orderbyquery.OrderByRowResult;
 import com.azure.cosmos.implementation.query.orderbyquery.OrderbyRowComparer;
-import com.azure.cosmos.models.ModelBridgeInternal;
 import reactor.core.publisher.Flux;
 
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
-import java.util.PriorityQueue;
-import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.PriorityBlockingQueue;
 import java.util.function.Function;
 
 public class NonStreamingOrderByUtils {
@@ -38,7 +39,7 @@ public static <T extends Resource> Flux<OrderByRowResult<Document>> nonStreaming
                 toNonStreamingOrderByQueryResultObservable(producer, tracker, queryMetricsMap, initialPageSize,
                     consumeComparer, clientSideRequestStatistics))
             .toArray(Flux[]::new);
-        return Flux.mergeOrdered(consumeComparer, fluxes);
+        return Flux.mergeComparingDelayError(1,consumeComparer, fluxes);
     }
 
     private static Flux<OrderByRowResult<Document>> toNonStreamingOrderByQueryResultObservable(DocumentProducer<Document> producer,
@@ -73,7 +74,7 @@ private PageToItemTransformer(RequestChargeTracker tracker, Map<String, QueryMet
 
         @Override
         public Flux<OrderByRowResult<Document>> apply(Flux<DocumentProducer<Document>.DocumentProducerFeedResponse> source) {
-            PriorityQueue<OrderByRowResult<Document>> priorityQueue = new PriorityQueue<>(consumeComparer);
+            PriorityBlockingQueue<OrderByRowResult<Document>> priorityQueue = new PriorityBlockingQueue<>(initialPageSize, consumeComparer);
 
             return source.flatMap(documentProducerFeedResponse -> {
                     clientSideRequestStatistics.addAll(
@@ -88,12 +89,17 @@ public Flux<OrderByRowResult<Document>> apply(Flux<DocumentProducer<Document>.Do
                             r.toJson(),
                             documentProducerFeedResponse.sourceFeedRange,
                             null);
-                        if (priorityQueue.size() < initialPageSize) {
-                            priorityQueue.add(orderByRowResult);
+                        if (Configs.getMaxItemSizeForVectorSearchEnabled()) {
+                            if (priorityQueue.size() < initialPageSize) {
+                                priorityQueue.add(orderByRowResult);
+                            } else {
+                                priorityQueue.add(orderByRowResult);
+                                priorityQueue.poll();
+                            }
                         } else {
                             priorityQueue.add(orderByRowResult);
-                            priorityQueue.poll();
                         }
+
                     });
                     tracker.addCharge(documentProducerFeedResponse.pageResult.getRequestCharge());
                     // Returning an empty Flux since we are only processing and managing state here

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

@@ -15,6 +15,8 @@
 
 import java.util.function.BiFunction;
 
+import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull;
+
 /**
  * While this class is public, but it is not part of our published public APIs.
  * This is meant to be internally used only by our sdk.
@@ -54,9 +56,10 @@ 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 ");
                     qryOptAccessor.getImpl(orderByCosmosQueryRequestOptions).setCustomItemSerializer(null);
                     documentQueryParams.setCosmosQueryRequestOptions(orderByCosmosQueryRequestOptions);
-                    return NonStreamingOrderByDocumentQueryExecutionContext.createAsync(diagnosticsClientContext, client, documentQueryParams, collection, 1000);
+                    return NonStreamingOrderByDocumentQueryExecutionContext.createAsync(diagnosticsClientContext, client, documentQueryParams, collection);
                 } else {
                     ModelBridgeInternal.setQueryRequestOptionsContinuationToken(orderByCosmosQueryRequestOptions, continuationToken);
                     qryOptAccessor.getImpl(orderByCosmosQueryRequestOptions).setCustomItemSerializer(null);

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

@@ -256,6 +256,26 @@ CosmosQueryRequestOptions setMaxItemCount(Integer maxItemCount) {
         return this;
     }
 
+    /**
+     * Gets the maximum item size to fetch during non-streaming order by queries.
+     *
+     * @return the max number of items for vector search.
+     */
+    public Integer getMaxItemSizeForVectorSearch() {
+        return this.actualRequestOptions.getMaxItemSizeForVectorSearch();
+    }
+
+    /**
+     * Sets the maximum item size to fetch during non-streaming order by queries.
+     *
+     * @param maxItemSizeForVectorSearch the max number of items for vector search.
+     * return the CosmosQueryRequestOptions.
+     */
+    public CosmosQueryRequestOptions setMaxItemSizeForVectorSearch(Integer maxItemSizeForVectorSearch) {
+        this.actualRequestOptions.setMaxItemSizeForVectorSearch(maxItemSizeForVectorSearch);
+        return this;
+    }
+
     /**
      * Gets the request continuation token.
      *

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

@@ -540,6 +540,11 @@ public static Integer getMaxItemCountFromQueryRequestOptions(CosmosQueryRequestO
         return options.getMaxItemCount();
     }
 
+    @Warning(value = INTERNAL_USE_ONLY_WARNING)
+    public static Integer getMaxItemSizeForVectorSearchFromQueryRequestOptions(CosmosQueryRequestOptions options) {
+        return options.getMaxItemSizeForVectorSearch();
+    }
+
     @Warning(value = INTERNAL_USE_ONLY_WARNING)
     public static String getRequestContinuationFromQueryRequestOptions(CosmosQueryRequestOptions options) {
         return options.getRequestContinuation();