diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/memo/Memo.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/memo/Memo.java index a225e0599c..a64db02847 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/memo/Memo.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/memo/Memo.java @@ -51,7 +51,6 @@ import java.util.Map; import java.util.Map.Entry; import java.util.Optional; import java.util.PriorityQueue; -import java.util.Queue; import java.util.stream.Collectors; import javax.annotation.Nullable; @@ -721,20 +720,24 @@ public class Memo { * In unrank() function, we will extract the actual physical function according the unique ID */ public long rank(long n) { + double threshold = 0.000000001; Preconditions.checkArgument(n > 0, "the n %d must be greater than 0 in nthPlan", n); List> plans = rankGroup(root, PhysicalProperties.GATHER); - Queue> rankingQueue = new PriorityQueue<>( - (l, r) -> -Double.compare(l.second, r.second)); - - for (Pair plan : plans) { - if (rankingQueue.size() == 0 || rankingQueue.size() < n) { - rankingQueue.add(plan); - } else if (rankingQueue.peek().second > plan.second) { - rankingQueue.poll(); - rankingQueue.add(plan); + plans = plans.stream().filter( + p -> !p.second.equals(Double.NaN) + && !p.second.equals(Double.POSITIVE_INFINITY) + && !p.second.equals(Double.NEGATIVE_INFINITY)) + .collect(Collectors.toList()); + // This is big heap, it always pops the element with larger cost or larger id. + PriorityQueue> pq = new PriorityQueue<>((l, r) -> Math.abs(l.second - r.second) < threshold + ? -Long.compare(l.first, r.first) : -Double.compare(l.second, r.second)); + for (Pair p : plans) { + pq.add(p); + if (pq.size() > n) { + pq.poll(); } } - return rankingQueue.peek().first; + return pq.peek().first; } private List> rankGroup(Group group, PhysicalProperties prop) {