[Cosmos][VectorSearch] Non Streaming Order By Query

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -4,6 +4,8 @@
 
 #### Features Added
 
+* 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
 
 #### Bugs Fixed

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 this.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,9 +4,11 @@
 
 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;
+import com.azure.cosmos.implementation.HttpConstants;
 import com.azure.cosmos.implementation.ImplementationBridgeHelpers;
 import com.azure.cosmos.implementation.OperationType;
 import com.azure.cosmos.implementation.PartitionKeyRange;
@@ -239,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) {
@@ -358,6 +361,37 @@ 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() && Configs.getMaxItemSizeForVectorSearchEnabled()) {
+            // Validate the TOP or LIMIT for non-streaming order-by queries
+            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()) {
+                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;
+                }
+            }
+        }
+
         // We need to compute the optimal initial page size for order-by queries
         if (queryInfo.hasOrderBy()) {
             int top;

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

@@ -0,0 +1,20 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+package com.azure.cosmos.implementation.query;
+
+import com.azure.cosmos.CosmosException;
+
+public class NonStreamingOrderByBadRequestException extends CosmosException {
+
+    private static final long serialVersionUID = 1L;
+
+    /**
+     * Creates a new instance of the NonStreamingOrderByBadRequestException class.
+     *
+     * @param statusCode the http status code of the response.
+     * @param errorMessage the error message.
+     */
+    public NonStreamingOrderByBadRequestException(int statusCode, String errorMessage) {
+        super(statusCode, errorMessage);
+    }
+}

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

@@ -0,0 +1,43 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.cosmos.implementation.query;
+
+import com.azure.cosmos.implementation.Document;
+import com.azure.cosmos.implementation.DocumentClientRetryPolicy;
+import com.azure.cosmos.implementation.RxDocumentServiceRequest;
+import com.azure.cosmos.implementation.feedranges.FeedRangeEpkImpl;
+import com.azure.cosmos.implementation.query.orderbyquery.OrderbyRowComparer;
+import com.azure.cosmos.models.CosmosQueryRequestOptions;
+import com.azure.cosmos.models.FeedResponse;
+import reactor.core.publisher.Mono;
+
+import java.util.UUID;
+import java.util.function.Function;
+import java.util.function.Supplier;
+
+public class NonStreamingOrderByDocumentProducer extends DocumentProducer<Document> {
+    private final OrderbyRowComparer<Document> consumeComparer;
+
+    NonStreamingOrderByDocumentProducer(
+        OrderbyRowComparer<Document> consumeComparer,
+        IDocumentQueryClient client,
+        String collectionResourceId,
+        CosmosQueryRequestOptions cosmosQueryRequestOptions,
+        TriFunction<FeedRangeEpkImpl, String, Integer, RxDocumentServiceRequest> createRequestFunc,
+        Function<RxDocumentServiceRequest, Mono<FeedResponse<Document>>> executeRequestFunc,
+        FeedRangeEpkImpl feedRange,
+        String collectionLink,
+        Supplier<DocumentClientRetryPolicy> createRetryPolicyFunc,
+        Class<Document> resourceType,
+        UUID correlatedActivityId,
+        int initialPageSize,
+        String initialContinuationToken,
+        int top,
+        Supplier<String> operationContextTextProvider) {
+        super(client, collectionResourceId, cosmosQueryRequestOptions, createRequestFunc, executeRequestFunc,
+            collectionLink, createRetryPolicyFunc, resourceType, correlatedActivityId, initialPageSize,
+            initialContinuationToken, top, feedRange, operationContextTextProvider);
+        this.consumeComparer = consumeComparer;
+    }
+}