add multi-orderby continuation token support (Azure#26267)

* add multi-orderby continuation token support

Co-authored-by: annie-mac <annie-mac@MININT-8V4CL0S.redmond.corp.microsoft.com>

Author: xinlian12

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

@@ -251,6 +251,7 @@ public static class HttpHeaders {
         public static final String QUERY_METRICS = "x-ms-documentdb-query-metrics";
         public static final String POPULATE_INDEX_METRICS = "x-ms-cosmos-populateindexmetrics";
         public static final String INDEX_UTILIZATION = "x-ms-cosmos-index-utilization";
+        public static final String QUERY_EXECUTION_INFO = "x-ms-cosmos-query-execution-info";
 
         // Batch operations
         public static final String IS_BATCH_ATOMIC = "x-ms-cosmos-batch-atomic";

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/WFConstants.java

@@ -76,5 +76,6 @@ public static class BackendHeaders {
         public static final String IS_USER_REQUEST = "x-ms-cosmos-internal-is-user-request";
         public static final String BACKEND_REQUEST_DURATION_MILLISECONDS = "x-ms-request-duration-ms";
         public static final String INDEX_UTILIZATION = "x-ms-cosmos-index-utilization";
+        public static final String QUERY_EXECUTION_INFO = "x-ms-cosmos-query-execution-info";
     }
 }

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdConstants.java

@@ -842,6 +842,7 @@ public enum RntbdResponseHeader implements RntbdHeader {
         IsRUPerMinuteUsed((short) 0x0027, RntbdTokenType.Byte, false),
         QueryMetrics((short) 0x0028, RntbdTokenType.String, false),
         IndexUtilization((short) 0x0044, RntbdTokenType.String, false),
+        QueryExecutionInfo((short) 0x0045, RntbdTokenType.String, false),
         GlobalCommittedLSN((short) 0x0029, RntbdTokenType.LongLong, false),
         NumberOfReadRegions((short) 0x0030, RntbdTokenType.ULong, false),
         OfferReplacePending((short) 0x0031, RntbdTokenType.Byte, false),

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdResponseHeaders.java

@@ -98,6 +98,8 @@ class RntbdResponseHeaders extends RntbdTokenStream<RntbdResponseHeader> {
     @JsonProperty
     private final RntbdToken indexUtilization;
     @JsonProperty
+    private final RntbdToken queryExecutionInfo;
+    @JsonProperty
     private final RntbdToken quorumAckedLSN;
     @JsonProperty
     private final RntbdToken quorumAckedLocalLSN;
@@ -174,6 +176,7 @@ private RntbdResponseHeaders(ByteBuf in) {
         this.queriesPerformed = this.get(RntbdResponseHeader.QueriesPerformed);
         this.queryMetrics = this.get(RntbdResponseHeader.QueryMetrics);
         this.indexUtilization = this.get(RntbdResponseHeader.IndexUtilization);
+        this.queryExecutionInfo = this.get(RntbdResponseHeader.QueryExecutionInfo);
         this.quorumAckedLSN = this.get(RntbdResponseHeader.QuorumAckedLSN);
         this.quorumAckedLocalLSN = this.get(RntbdResponseHeader.QuorumAckedLocalLSN);
         this.readsPerformed = this.get(RntbdResponseHeader.ReadsPerformed);
@@ -276,6 +279,7 @@ public void setValues(final Map<String, String> headers) {
         this.mapValue(this.partitionKeyRangeId, BackendHeaders.PARTITION_KEY_RANGE_ID, String::toString, headers);
         this.mapValue(this.queryMetrics, BackendHeaders.QUERY_METRICS, String::toString, headers);
         this.mapValue(this.indexUtilization, BackendHeaders.INDEX_UTILIZATION, String::toString, headers);
+        this.mapValue(this.queryExecutionInfo, BackendHeaders.QUERY_EXECUTION_INFO, String::toString, headers);
         this.mapValue(this.quorumAckedLSN, BackendHeaders.QUORUM_ACKED_LSN, Long::parseLong, headers);
         this.mapValue(this.quorumAckedLocalLSN, BackendHeaders.QUORUM_ACKED_LOCAL_LSN, Long::parseLong, headers);
         this.mapValue(this.requestCharge, HttpHeaders.REQUEST_CHARGE, Double::parseDouble, headers);
@@ -422,6 +426,9 @@ private void collectEntries(final BiConsumer<RntbdToken, Function<RntbdToken, Ma
             toStringEntry(BackendHeaders.INDEX_UTILIZATION, token)
         );
 
+        collector.accept(this.queryExecutionInfo, tokens ->
+            toStringEntry(BackendHeaders.QUERY_EXECUTION_INFO, tokens));
+
         collector.accept(this.quorumAckedLSN, token ->
             toLongEntry(BackendHeaders.QUORUM_ACKED_LSN, token)
         );

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

@@ -19,10 +19,12 @@
 import com.azure.cosmos.implementation.Undefined;
 import com.azure.cosmos.implementation.Utils;
 import com.azure.cosmos.implementation.Utils.ValueHolder;
-import com.azure.cosmos.implementation.apachecommons.lang.NotImplementedException;
 import com.azure.cosmos.implementation.apachecommons.lang.tuple.ImmutablePair;
 import com.azure.cosmos.implementation.apachecommons.lang.tuple.Pair;
 import com.azure.cosmos.implementation.feedranges.FeedRangeEpkImpl;
+import com.azure.cosmos.implementation.query.orderbyquery.ComparisonFilters;
+import com.azure.cosmos.implementation.query.orderbyquery.ComparisonWithDefinedFilters;
+import com.azure.cosmos.implementation.query.orderbyquery.ComparisonWithUndefinedFilters;
 import com.azure.cosmos.implementation.query.orderbyquery.OrderByRowResult;
 import com.azure.cosmos.implementation.query.orderbyquery.OrderbyRowComparer;
 import com.azure.cosmos.implementation.routing.Range;
@@ -279,47 +281,26 @@ private OrderByDocumentQueryExecutionContext<T>.FormattedFilterInfo getFormatted
             SortOrder sortOrder = sortOrders[0];
             QueryItem orderByItem = orderByItems[0];
             Object rawItem = orderByItem.getItem();
-            String orderByItemToString;
-            if (rawItem instanceof String) {
-                orderByItemToString = "\"" + rawItem.toString().replaceAll("\"",
-                        "\\\"") + "\"";
-            } else {
-                if (rawItem != null) {
-                    orderByItemToString = rawItem.toString();
-                } else {
-                    orderByItemToString = "null";
-                }
-            }
-            if (rawItem == Undefined.value()) {
-                // Handling undefined needs filter literals
-                // What we really want is to support expression > undefined,
-                // but the engine evaluates to undefined instead of true or false,
-                // so we work around this by using the IS_DEFINED() system function.
 
-                left.append(sortOrder == SortOrder.Descending ? "false" : "IS_DEFINED(" + expression + ") ");
-                target.append(sortOrder == SortOrder.Descending ? "NOT IS_DEFINED(" + expression + ")" : "true ");
-                right.append(sortOrder == SortOrder.Descending ? "NOT IS_DEFINED(" + expression + ")" : "true ");
+            this.appendToBuilders(left, target, right, "(");
 
-            } else {
+            String orderByItemToString = this.getOrderByItemString(rawItem);
 
-                left.append(getFilterString(expression,
-                                            (sortOrder == SortOrder.Descending ? "<" : ">"),
-                                            orderByItemToString));
-
-                if (inclusive) {
-                    target.append(getFilterString(expression,
-                                                  (sortOrder == SortOrder.Descending ? "<=" : ">="),
-                                                  orderByItemToString));
-                } else {
-                    target.append(getFilterString(expression,
-                                                  (sortOrder == SortOrder.Descending ? "<" : ">"),
-                                                  orderByItemToString));
-                }
+             // Handling undefined needs filter literals
+             // What we really want is to support expression > undefined,
+             // but the engine evaluates to undefined instead of true or false,
+             // so we work around this by using the IS_DEFINED() system function
+             // ComparisonWithUndefinedFilters is used to handle the logic mentioned above
+            ComparisonFilters filters =
+                rawItem == Undefined.value() ? new ComparisonWithUndefinedFilters(expression) : new ComparisonWithDefinedFilters(expression, orderByItemToString);
 
-                right.append(getFilterString(expression,
-                                             (sortOrder == SortOrder.Descending ? "<=" : ">="),
-                                             orderByItemToString));
+            left.append(sortOrder == SortOrder.Descending ? filters.lessThan() : filters.greaterThan());
+            if (inclusive) {
+                target.append(sortOrder == SortOrder.Descending ? filters.lessThanOrEqualTo() : filters.greaterThanOrEqualTo());
+            } else {
+                target.append(sortOrder == SortOrder.Descending ? filters.lessThan() : filters.greaterThan());
             }
+            right.append(sortOrder == SortOrder.Descending ? filters.lessThanOrEqualTo() : filters.greaterThanOrEqualTo());
 
             // Now we need to include all the types that match the sort order.
             List<String> definedFunctions =
@@ -336,20 +317,148 @@ private OrderByDocumentQueryExecutionContext<T>.FormattedFilterInfo getFormatted
             target.append(isDefinedFunctions);
             right.append(isDefinedFunctions);
 
+            this.appendToBuilders(left, target, right, ")");
+
         } else {
-            // This code path needs to be implemented, but it's error prone and needs
-            // testing.
-            // You can port the implementation from the .net SDK and it should work if
-            // ported right.
-            throw new NotImplementedException(
-                    "Resuming a multi order by query from a continuation token is not supported yet.");
+            // For a multi order by query
+            // Suppose the query is SELECT* FROM c ORDER BY c.string ASC, c.number ASC
+            // And we left off on partition N with the value("A", 1)
+            // Then
+            //      All the partitions to the left will have finished reading("A", 1)
+            //      Partition N is still reading("A", 1)
+            //      All the partitions to the right have let to read a "(A", 1)
+            // The filters are harder to derive since there are multiple columns
+            // But the problem reduces to "How do you know one document comes after another in a multi order by query"
+            // The answer is to just look at it one column at a time.
+            // For this particular scenario:
+            //      If a first column is greater ex. ("B", blah), then the document comes later in the sort order
+            //      Therefore we want all documents where the first column is greater than "A" which means > "A"
+            //      Or if the first column is a tie, then you look at the second column ex. ("A", blah).
+            //      Therefore we also want all documents where the first column was a tie but the second column is greater which means = "A" AND > 1
+            //      Therefore the filters should be
+            //      (> "A") OR (= "A" AND > 1), (> "A") OR (= "A" AND >= 1), (> "A") OR (= "A" AND >= 1)
+            //      Notice that if we repeated the same logic we for single order by we would have gotten
+            //      > "A" AND > 1, >= "A" AND >= 1, >= "A" AND >= 1
+            //      which is wrong since we missed some documents
+            //      Repeat the same logic for ASC, DESC
+            //          (> "A") OR (= "A" AND < 1), (> "A") OR (= "A" AND <= 1), (> "A") OR (= "A" AND <= 1)
+            //      Again for DESC, ASC
+            //          (< "A") OR (= "A" AND > 1), (< "A") OR (= "A" AND >= 1), (< "A") OR (= "A" AND >= 1)
+            //      And again for DESC DESC
+            //          (< "A") OR (= "A" AND < 1), (< "A") OR (= "A" AND <= 1), (< "A") OR (= "A" AND <= 1)
+            //      The general we look at all prefixes of the order by columns to look for tie breakers.
+            //      Except for the full prefix whose last column follows the rules for single item order by
+            //      And then you just OR all the possibilities together
+
+            for (int prefixLength = 1; prefixLength <= numOrderByItems; prefixLength++) {
+                boolean lastPrefix = prefixLength == numOrderByItems;
+
+                this.appendToBuilders(left, target, right, "(");
+
+                for (int index = 0; index < prefixLength; index++) {
+                    String expression = expressions[index];
+                    SortOrder sortOrder = sortOrders[index];
+                    QueryItem orderbyItem = orderByItems[index];
+                    Object orderbyRawItem = orderbyItem.getItem();
+
+                    boolean lastItem = index == prefixLength - 1;
+
+                    this.appendToBuilders(left, target, right, "(");
+                    String orderByItemToString = getOrderByItemString(orderbyRawItem);
+                    ComparisonFilters filters =
+                        orderbyRawItem == Undefined.value() ? new ComparisonWithUndefinedFilters((expression)) : new ComparisonWithDefinedFilters(expression, orderByItemToString);
+
+                    if (lastItem) {
+                        if (lastPrefix) {
+                            left.append(sortOrder == SortOrder.Descending ? filters.lessThan() : filters.greaterThan());
+
+                            if (inclusive) {
+                                target.append(sortOrder == SortOrder.Descending ? filters.lessThanOrEqualTo() : filters.greaterThanOrEqualTo());
+                            } else {
+                                target.append(sortOrder == SortOrder.Descending ? filters.lessThan() : filters.greaterThan());
+                            }
+
+                            right.append(sortOrder == SortOrder.Descending ? filters.lessThanOrEqualTo() : filters.greaterThanOrEqualTo());
+                        } else {
+                            left.append(sortOrder == SortOrder.Descending ? filters.lessThan() : filters.greaterThan());
+                            target.append(sortOrder == SortOrder.Descending ? filters.lessThan() : filters.greaterThan());
+                            right.append(sortOrder == SortOrder.Descending ? filters.lessThan() : filters.greaterThan());
+                        }
+
+                    } else {
+                        left.append(filters.equalTo());
+                        target.append(filters.equalTo());
+                        right.append(filters.equalTo());
+                    }
+
+                    if (lastItem) {
+                        // Now we need to include all the types that match the sort order.
+                        List<String> definedFunctions =
+                            IsSystemFunctions.getIsDefinedFunctions(ItemTypeHelper.getOrderByItemType(orderbyRawItem),
+                                sortOrder == SortOrder.Ascending);
+                        StringBuilder isDefinedFuncBuilder = new StringBuilder();
+                        for (String idf : definedFunctions) {
+                            isDefinedFuncBuilder.append(" OR ");
+                            isDefinedFuncBuilder.append(String.format("%s(%s)", idf, expression));
+                        }
+
+                        String isDefinedFunctions = isDefinedFuncBuilder.toString();
+                        left.append(isDefinedFunctions);
+                        target.append(isDefinedFunctions);
+                        right.append(isDefinedFunctions);
+                    }
+
+                    this.appendToBuilders(left, target, right, ")");
+                    if (!lastItem) {
+                        this.appendToBuilders(left, target, right, " AND ");
+                    }
+                }
+
+                this.appendToBuilders(left, target, right, ")");
+                if (!lastPrefix) {
+                    this.appendToBuilders(left, target, right, " OR ");
+                }
+            }
         }
 
         return new FormattedFilterInfo(left.toString(),
                 target.toString(),
                 right.toString());
     }
 
+    private String getOrderByItemString(Object orderbyRawItem) {
+        String orderByItemToString;
+        if (orderbyRawItem instanceof String) {
+            orderByItemToString = "\"" + orderbyRawItem.toString().replaceAll("\"",
+                "\\\"") + "\"";
+        } else {
+            if (orderbyRawItem != null) {
+                orderByItemToString = orderbyRawItem.toString();
+            } else {
+                orderByItemToString = "null";
+            }
+        }
+
+        return orderByItemToString;
+    }
+
+    private void appendToBuilders(StringBuilder leftBuilder, StringBuilder targetBuilder, StringBuilder rightBuilder, String appendText) {
+        this.appendToBuilders(leftBuilder, targetBuilder, rightBuilder, appendText, appendText, appendText);
+    }
+
+    private void appendToBuilders(
+        StringBuilder leftBuilder,
+        StringBuilder targetBuilder,
+        StringBuilder rightBuilder,
+        String leftAppendText,
+        String targetAppendText,
+        String rightAppendText) {
+
+        leftBuilder.append(leftAppendText);
+        targetBuilder.append(targetAppendText);
+        rightBuilder.append(rightAppendText);
+    }
+
     static class IsSystemFunctions {
         static final String Defined = "IS_DEFINED";
         static final String NotDefined = "NOT IS_DEFINED";
@@ -420,10 +529,6 @@ List<String> getExtendedTypesIsDefinedFunctions(int index, boolean isAscending)
         }
     }
 
-    private String getFilterString(String s1, String s2, String s3) {
-        return String.format("%s %s %s", s1, s2, s3);
-    }
-
     @Override
     protected OrderByDocumentProducer<T> createDocumentProducer(
             String collectionRid,