///|
/// One scored document returned by a collector.
pub(all) struct SearchHit {
address : DocAddress
score : Double
} derive(@debug.Debug)
///|
/// Collector that retains at most the highest-scoring K documents.
pub struct TopKCollector {
limit : Int
}
///|
pub fn TopKCollector::new(limit : Int) -> TopKCollector {
{ limit, }
}
///|
fn compare_hits(left : SearchHit, right : SearchHit) -> Int {
if left.score > right.score {
-1
} else if left.score < right.score {
1
} else if left.address.segment_ord < right.address.segment_ord {
-1
} else if left.address.segment_ord > right.address.segment_ord {
1
} else {
left.address.doc_id.value.compare(right.address.doc_id.value)
}
}
///|
fn TopKCollector::collect_into(
self : TopKCollector,
scorer : &Scorer,
segment_ord : Int,
deleted_docs : ReadOnlyArray[DocId],
candidates : Array[SearchHit],
) -> Unit {
if self.limit <= 0 {
return
}
while scorer.advance() {
let doc_id = scorer.doc()
if deleted_docs.search_by(deleted => deleted == doc_id) is None {
candidates.push({
address: DocAddress::new(segment_ord, doc_id),
score: scorer.score(),
})
}
}
}
///|
fn TopKCollector::finish(
self : TopKCollector,
candidates : Array[SearchHit],
) -> ReadOnlyArray[SearchHit] {
if self.limit <= 0 {
return []
}
candidates.sort_by(compare_hits)
let result : Array[SearchHit] = []
let result_count = if candidates.length() < self.limit {
candidates.length()
} else {
self.limit
}
for index in 0..