rubenada commented on a change in pull request #2109:
URL: https://github.com/apache/calcite/pull/2109#discussion_r469942516
##########
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);
+ }
+ if (l.size() == 1) {
+ l = new ArrayList<>(l);
+ }
+ l.add(o);
+ return l;
+ });
+ size++;
+ }
+ }
+
+ if (offset > 0) {
+ // search until which key we have to remove entries from the map
+ int skipped = 0;
+ TKey until = null;
+ for (Map.Entry<TKey, List<TSource>> e : map.entrySet()) {
+ skipped += e.getValue().size();
+
+ if (skipped > offset) {
+ // we might need to remove entries from the list
+ List<TSource> l = e.getValue();
+ int toKeep = skipped - offset;
+ if (toKeep < l.size()) {
+ l.subList(0, l.size() - toKeep).clear();
+ }
+
+ until = e.getKey();
+ break;
+ }
+ }
+ if (until == null) {
+ return Linq4j.emptyEnumerator();
+ } else if (until != null) {
Review comment:
this if condition looks wrong/useless
----------------------------------------------------------------
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]
