///|
pub enum RunOrder {
ScoreDescending
DocumentAscending
InputOrder
} derive(Debug, Eq, ToJson)
///|
pub fn RunOrder::score_descending() -> RunOrder {
ScoreDescending
}
///|
pub fn RunOrder::document_ascending() -> RunOrder {
DocumentAscending
}
///|
pub fn RunOrder::input_order() -> RunOrder {
InputOrder
}
///|
pub fn sort_run(
run : Array[RetrievedDoc],
order : RunOrder,
) -> Array[RetrievedDoc] {
let result = run.copy()
match order {
ScoreDescending =>
result.sort_by(fn(a, b) {
let score_order = b.score.compare(a.score)
if score_order == 0 {
a.doc_id.compare(b.doc_id)
} else {
score_order
}
})
DocumentAscending =>
result.sort_by(fn(a, b) {
let query_order = a.query_id.compare(b.query_id)
if query_order == 0 {
a.doc_id.compare(b.doc_id)
} else {
query_order
}
})
InputOrder => ()
}
result
}
///|
pub fn truncate_run(
run : Array[RetrievedDoc],
per_query : Int,
) -> Array[RetrievedDoc] {
let limit = Int::max(per_query, 0)
let grouped = group_runs(run)
let result : Array[RetrievedDoc] = []
let ids : Array[String] = []
for query_id, _ in grouped {
ids.push(query_id)
}
ids.sort()
for query_id in ids {
let ranked = grouped[query_id]
for index in 0.. Array[RetrievedDoc] {
run.filter(fn(item) { item.score >= minimum_score })
}
///|
pub fn merge_runs(
left : Array[RetrievedDoc],
right : Array[RetrievedDoc],
) -> Array[RetrievedDoc] {
let best : Map[String, RetrievedDoc] = Map([])
for item in left {
best["\{item.query_id}:\{item.doc_id}"] = item
}
for item in right {
let key = "\{item.query_id}:\{item.doc_id}"
match best.get(key) {
Some(existing) => if item.score > existing.score { best[key] = item }
None => best[key] = item
}
}
let result : Array[RetrievedDoc] = []
for _, item in best {
result.push(item)
}
sort_run(result, ScoreDescending)
}
///|
pub fn run_to_pool(
run : Array[RetrievedDoc],
query_id : String,
) -> CandidatePool {
let ranked = group_runs(run).get_or_default(query_id, [])
let docs : Array[String] = []
let seen : Map[String, Unit] = Map([])
for item in ranked {
if !seen.contains(item.doc_id) {
seen[item.doc_id] = ()
docs.push(item.doc_id)
}
}
{ query_id, doc_ids: docs }
}
///|
pub fn normalize_scores(run : Array[RetrievedDoc]) -> Array[RetrievedDoc] {
let grouped = group_runs(run)
let result : Array[RetrievedDoc] = []
for _, ranked in grouped {
if ranked.is_empty() {
continue
}
let mut minimum = ranked[0].score
let mut maximum = ranked[0].score
for item in ranked {
minimum = Double::min(minimum, item.score)
maximum = Double::max(maximum, item.score)
}
let width = maximum - minimum
for item in ranked {
let score = if width == 0.0 {
1.0
} else {
(item.score - minimum) / width
}
result.push({ query_id: item.query_id, doc_id: item.doc_id, score })
}
}
sort_run(result, ScoreDescending)
}
///|
pub fn blend_runs(
left : Array[RetrievedDoc],
right : Array[RetrievedDoc],
left_weight? : Double = 0.5,
) -> Array[RetrievedDoc] {
let weight = Double::min(Double::max(left_weight, 0.0), 1.0)
let left_scores : Map[String, Double] = Map([])
let right_scores : Map[String, Double] = Map([])
for item in normalize_scores(left) {
left_scores["\{item.query_id}:\{item.doc_id}"] = item.score
}
for item in normalize_scores(right) {
right_scores["\{item.query_id}:\{item.doc_id}"] = item.score
}
let keys : Map[String, Unit] = Map([])
for key, _ in left_scores {
keys[key] = ()
}
for key, _ in right_scores {
keys[key] = ()
}
let result : Array[RetrievedDoc] = []
for key, _ in keys {
let pieces = key.split(":").to_array()
if pieces.length() >= 2 {
let query_id = pieces[0].to_owned()
let doc_id = pieces[1].to_owned()
let score = weight * left_scores.get_or_default(key, 0.0) +
(1.0 - weight) * right_scores.get_or_default(key, 0.0)
result.push({ query_id, doc_id, score })
}
}
sort_run(result, ScoreDescending)
}