This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new ad4e32bf66 [global-index] Optimize hybrid ranker top-k (#8348)
ad4e32bf66 is described below
commit ad4e32bf66b0e4e93f78381ef3124130f3ddde98
Author: QuakeWang <[email protected]>
AuthorDate: Thu Jun 25 10:31:06 2026 +0800
[global-index] Optimize hybrid ranker top-k (#8348)
Java hybrid ranker currently sorts all fused candidates before
truncating to topK, making the final selection O(n log n). This PR
replaces that full sort with a bounded heap, reducing it to O(n log
limit) while preserving the existing score desc / rowId asc tie-breaking
semantics.
A regression test is added to cover boundary ties, ensuring smaller
rowId candidates are retained when scores are equal.
---
.../paimon/globalindex/HybridSearchRanker.java | 38 +++++++++++++++++-----
.../paimon/globalindex/HybridSearchRankerTest.java | 20 ++++++++++++
2 files changed, 49 insertions(+), 9 deletions(-)
diff --git
a/paimon-common/src/main/java/org/apache/paimon/globalindex/HybridSearchRanker.java
b/paimon-common/src/main/java/org/apache/paimon/globalindex/HybridSearchRanker.java
index 1b6504c0f9..675e4b1a38 100644
---
a/paimon-common/src/main/java/org/apache/paimon/globalindex/HybridSearchRanker.java
+++
b/paimon-common/src/main/java/org/apache/paimon/globalindex/HybridSearchRanker.java
@@ -22,9 +22,11 @@ import org.apache.paimon.utils.RoaringNavigableMap64;
import java.io.Serializable;
import java.util.ArrayList;
+import java.util.Comparator;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
+import java.util.PriorityQueue;
/** Ranker utilities for hybrid search results. */
public class HybridSearchRanker {
@@ -151,24 +153,42 @@ public class HybridSearchRanker {
}
private static ScoredGlobalIndexResult topK(Map<Long, Float> scores, int
limit) {
- if (scores.isEmpty()) {
+ if (scores.isEmpty() || limit <= 0) {
return ScoredGlobalIndexResult.createEmpty();
}
- List<Map.Entry<Long, Float>> ranked = new
ArrayList<>(scores.entrySet());
- ranked.sort(
+ if (scores.size() <= limit) {
+ return toScoredResult(scores.entrySet());
+ }
+
+ // The heap head is the weakest kept row: lowest score, then largest
rowId.
+ Comparator<Map.Entry<Long, Float>> weakestFirst =
(left, right) -> {
- int scoreCompare = Float.compare(right.getValue(),
left.getValue());
+ int scoreCompare = Float.compare(left.getValue(),
right.getValue());
if (scoreCompare != 0) {
return scoreCompare;
}
- return Long.compare(left.getKey(), right.getKey());
- });
+ return Long.compare(right.getKey(), left.getKey());
+ };
+
+ PriorityQueue<Map.Entry<Long, Float>> topEntries =
+ new PriorityQueue<>(limit + 1, weakestFirst);
+ for (Map.Entry<Long, Float> entry : scores.entrySet()) {
+ if (topEntries.size() < limit) {
+ topEntries.offer(entry);
+ } else if (weakestFirst.compare(entry, topEntries.peek()) > 0) {
+ topEntries.poll();
+ topEntries.offer(entry);
+ }
+ }
+
+ return toScoredResult(topEntries);
+ }
- int size = Math.min(limit, ranked.size());
+ private static ScoredGlobalIndexResult toScoredResult(
+ Iterable<Map.Entry<Long, Float>> entries) {
RoaringNavigableMap64 bitmap = new RoaringNavigableMap64();
Map<Long, Float> topScores = new HashMap<>();
- for (int i = 0; i < size; i++) {
- Map.Entry<Long, Float> entry = ranked.get(i);
+ for (Map.Entry<Long, Float> entry : entries) {
bitmap.add(entry.getKey());
topScores.put(entry.getKey(), entry.getValue());
}
diff --git
a/paimon-common/src/test/java/org/apache/paimon/globalindex/HybridSearchRankerTest.java
b/paimon-common/src/test/java/org/apache/paimon/globalindex/HybridSearchRankerTest.java
index 0d4b2c0824..0e3ad8cac5 100644
---
a/paimon-common/src/test/java/org/apache/paimon/globalindex/HybridSearchRankerTest.java
+++
b/paimon-common/src/test/java/org/apache/paimon/globalindex/HybridSearchRankerTest.java
@@ -127,6 +127,26 @@ public class HybridSearchRankerTest {
assertThat(ranked.scoreGetter().score(3L)).isCloseTo(2.0f,
within(0.000001f));
}
+ @Test
+ public void testWeightedScoreTopKBreaksBoundaryTiesByRowId() {
+ ScoredGlobalIndexResult tied =
+ result(
+ new long[] {4, 3, 2, 1},
+ new float[] {5.0f, 5.0f, 5.0f, 5.0f},
+ new long[] {4, 3, 2, 1});
+
+ ScoredGlobalIndexResult ranked =
+ HybridSearchRanker.weightedScore(
+ Collections.singletonList(
+ new HybridSearchRanker.WeightedResult(tied,
2.0f)),
+ 2);
+
+ assertThat(ranked.results()).contains(1L, 2L);
+ assertThat(ranked.results()).doesNotContain(3L, 4L);
+ assertThat(ranked.scoreGetter().score(1L)).isCloseTo(2.0f,
within(0.000001f));
+ assertThat(ranked.scoreGetter().score(2L)).isCloseTo(2.0f,
within(0.000001f));
+ }
+
@Test
public void testRejectNonFiniteWeights() {
ScoredGlobalIndexResult result = result(new long[] {1}, new float[]
{1.0f});
