thomasrebele commented on a change in pull request #2109:
URL: https://github.com/apache/calcite/pull/2109#discussion_r472859765



##########
File path: 
linq4j/src/main/java/org/apache/calcite/linq4j/EnumerableDefaults.java
##########
@@ -2624,6 +2624,86 @@ public static boolean isMergeJoinSupported(JoinType 
joinType) {
     };
   }
 
+
+  /**
+   * A sort implementation optimized for a sort with a fetch size (LIMIT).
+   * @param offset how many rows are skipped from the sorted output.
+   *               Must be greater than or equal to 0.
+   * @param fetch how many rows are retrieved. Must be greater than 0.
+   */
+  public static <TSource, TKey> Enumerable<TSource> orderBy(
+      Enumerable<TSource> source,
+      Function1<TSource, TKey> keySelector,
+      Comparator<TKey> comparator,
+      int offset, int fetch) {
+    return new AbstractEnumerable<TSource>() {
+      @Override public Enumerator<TSource> enumerator() {
+        TreeMap<TKey, List<TSource>> map = new TreeMap<>(comparator);
+        long size = 0;
+        long needed = fetch + offset;
+
+        try (Enumerator<TSource> os = source.enumerator()) {
+          while (os.moveNext()) {
+            TSource o = os.current();
+            TKey key = keySelector.apply(o);
+            if (needed >= 0 && size >= needed) {
+              if (comparator.compare(key, map.lastKey()) >= 0) {
+                continue;
+              }
+              // remove last entry from tree map
+              List<TSource> l = map.get(map.lastKey());
+              if (l.size() == 1) {
+                map.remove(map.lastKey());
+              } else {
+                l.remove(l.size() - 1);
+              }
+              size--;
+            }
+            map.compute(key, (k, l) -> {
+              if (l == null) {
+                return Collections.singletonList(o);

Review comment:
       I've added an explanation at the beginning and more comments in the 
actual code to make the algorithm easier to understand.




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
[email protected]


Reply via email to