Skip to content

Commit 8b2911b

Browse files
authored
[core] Stabilize scored global index top-k tie-breaking (#8301)
`ScoredGlobalIndexResult.topK` used a min-heap ordered only by score, and replaced the heap head only when a new row had a strictly higher score. When rows tied around the top-k boundary, the retained row IDs could differ from the expected `score desc, rowId asc` ranking semantics.
1 parent 23fca3c commit 8b2911b

2 files changed

Lines changed: 88 additions & 6 deletions

File tree

paimon-common/src/main/java/org/apache/paimon/globalindex/ScoredGlobalIndexResult.java

Lines changed: 11 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -85,17 +85,22 @@ default ScoredGlobalIndexResult topK(int k) {
8585
}
8686

8787
ScoreGetter scoreGetter = scoreGetter();
88-
// Min-heap by score: the head is the smallest score so we can evict it when a
89-
// higher-scored row arrives. This gives O(n log k) instead of O(n log n).
90-
PriorityQueue<long[]> minHeap =
91-
new PriorityQueue<>(
92-
k + 1, Comparator.comparingDouble(a -> Float.intBitsToFloat((int) a[1])));
88+
// Min-heap whose ordering matches the global index ranking semantics (score desc,
89+
// rowId asc): the head is the weakest candidate currently kept, i.e. the lowest
90+
// score and, among ties, the largest rowId. A new row replaces the head only when
91+
// it is strictly stronger, so the retained set equals a full "score desc, rowId asc"
92+
// sort truncated to k, while keeping O(n log k) instead of O(n log n).
93+
// entry: [rowId, rawScoreBits]
94+
Comparator<long[]> weakestFirst =
95+
Comparator.<long[]>comparingDouble(a -> Float.intBitsToFloat((int) a[1]))
96+
.thenComparing(Comparator.comparingLong((long[] a) -> a[0]).reversed());
97+
PriorityQueue<long[]> minHeap = new PriorityQueue<>(k + 1, weakestFirst);
9398
for (long rowId : rowIds) {
9499
float score = scoreGetter.score(rowId);
95100
long[] entry = new long[] {rowId, Float.floatToRawIntBits(score)};
96101
if (minHeap.size() < k) {
97102
minHeap.offer(entry);
98-
} else if (score > Float.intBitsToFloat((int) minHeap.peek()[1])) {
103+
} else if (weakestFirst.compare(entry, minHeap.peek()) > 0) {
99104
minHeap.poll();
100105
minHeap.offer(entry);
101106
}
Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,77 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing, software
13+
* distributed under the License is distributed on an "AS IS" BASIS,
14+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15+
* See the License for the specific language governing permissions and
16+
* limitations under the License.
17+
*/
18+
19+
package org.apache.paimon.globalindex;
20+
21+
import org.apache.paimon.utils.RoaringNavigableMap64;
22+
23+
import org.junit.jupiter.api.Test;
24+
25+
import java.util.HashMap;
26+
import java.util.Map;
27+
28+
import static org.assertj.core.api.Assertions.assertThat;
29+
30+
/** Tests for {@link ScoredGlobalIndexResult}. */
31+
public class ScoredGlobalIndexResultTest {
32+
33+
@Test
34+
public void testTopKBreaksBoundaryTiesByRowId() {
35+
ScoredGlobalIndexResult result =
36+
result(new long[] {1, 2, 3, 4}, new float[] {0.5f, 0.5f, 0.5f, 0.9f});
37+
38+
RoaringNavigableMap64 topK = result.topK(2).results();
39+
40+
assertThat(topK.getIntCardinality()).isEqualTo(2);
41+
assertThat(topK).contains(1L, 4L);
42+
assertThat(topK).doesNotContain(2L, 3L);
43+
}
44+
45+
@Test
46+
public void testTopKBreaksTieGroupBySmallerRowId() {
47+
ScoredGlobalIndexResult result =
48+
result(new long[] {1, 2, 3, 4, 5}, new float[] {0.5f, 0.5f, 0.5f, 0.5f, 0.9f});
49+
50+
// The high-scored row 5 forces a heap eviction within the 0.5 tie group, so this
51+
// distinguishes the score-only heap ({2, 3, 5}) from the fixed heap ({1, 2, 5}).
52+
RoaringNavigableMap64 topK = result.topK(3).results();
53+
54+
assertThat(topK.getIntCardinality()).isEqualTo(3);
55+
assertThat(topK).contains(1L, 2L, 5L);
56+
assertThat(topK).doesNotContain(3L, 4L);
57+
}
58+
59+
@Test
60+
public void testTopKReturnsSameResultWhenCardinalityDoesNotExceedK() {
61+
ScoredGlobalIndexResult result =
62+
result(new long[] {1, 2, 3}, new float[] {0.1f, 0.2f, 0.3f});
63+
64+
assertThat(result.topK(3)).isSameAs(result);
65+
assertThat(result.topK(5)).isSameAs(result);
66+
}
67+
68+
private ScoredGlobalIndexResult result(long[] rowIds, float[] scores) {
69+
RoaringNavigableMap64 bitmap = new RoaringNavigableMap64();
70+
Map<Long, Float> scoreMap = new HashMap<>();
71+
for (int i = 0; i < rowIds.length; i++) {
72+
bitmap.add(rowIds[i]);
73+
scoreMap.put(rowIds[i], scores[i]);
74+
}
75+
return ScoredGlobalIndexResult.create(bitmap, scoreMap::get);
76+
}
77+
}

0 commit comments

Comments
 (0)