Skip to content

Commit 4c1418e

Browse files
fix(index): sort vector search top-k results when not truncating
`SearchResult::top_k` only sorted results by score on the truncation path (candidates > k). When `k >= candidate count`, it returned rows in the unordered scored-map / insertion order instead of by relevance. Because the row order was previously discarded downstream (row-range scans re-read in file order), this went unnoticed — but consumers that honor the row order (e.g. the DataFusion `vector_search` scan) then get unranked results whenever `limit >= number of matched rows`. This affects `vector_search`, `full_text_search`, and `hybrid_search`. Sort the no-truncation branch by relevance rank too, and update the unit tests that asserted the old insertion-order behavior.
1 parent 27dea79 commit 4c1418e

1 file changed

Lines changed: 13 additions & 11 deletions

File tree

crates/paimon/src/vector_search.rs

Lines changed: 13 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -218,12 +218,11 @@ impl SearchResult {
218218
}
219219

220220
if best_by_row_id.len() <= k {
221-
// Keep the original row order when no truncation is needed.
222-
let rows = self
223-
.row_ids
224-
.iter()
225-
.filter_map(|row_id| best_by_row_id.remove(row_id))
226-
.collect();
221+
// No truncation needed, but still return best-first: the scored map is
222+
// unordered, and consumers of the row order (e.g. the DataFusion
223+
// vector_search scan) require relevance rank, not insertion order.
224+
let mut rows: Vec<ScoredRow> = best_by_row_id.into_values().collect();
225+
sort_scored_rows_by_rank(&mut rows);
227226
return Self::from_scored_rows(rows);
228227
}
229228

@@ -382,12 +381,14 @@ mod tests {
382381
}
383382

384383
#[test]
385-
fn test_search_result_top_k_preserves_order_without_truncation() {
384+
fn test_search_result_top_k_sorts_best_first_without_truncation() {
385+
// Even when k >= candidate count (no truncation), results must be returned
386+
// best-first by score, not in the input/insertion order.
386387
let result = SearchResult::new(vec![3, 1, 2], vec![0.1, 0.9, 0.5]);
387388

388389
let top = result.top_k(3);
389-
assert_eq!(top.row_ids, result.row_ids);
390-
assert_eq!(top.scores, result.scores);
390+
assert_eq!(top.row_ids, vec![1, 2, 3]);
391+
assert_eq!(top.scores, vec![0.9, 0.5, 0.1]);
391392
}
392393

393394
#[test]
@@ -409,8 +410,9 @@ mod tests {
409410
.without_deleted_row_ranges(Some(&deleted))
410411
.unwrap()
411412
.top_k(10);
412-
assert_eq!(filtered.row_ids, vec![1, 4]);
413-
assert_eq!(filtered.scores, vec![0.1, 0.2]);
413+
// Rows 2..3 are deleted; the remaining rows come back best-first by score.
414+
assert_eq!(filtered.row_ids, vec![4, 1]);
415+
assert_eq!(filtered.scores, vec![0.2, 0.1]);
414416
}
415417

416418
#[test]

0 commit comments

Comments
 (0)