feat: return BM25 scores for FTS (#2687)

We need this for reranking Signed-off-by: BubbleCal <bubble-cal@outlook.com>
lancedb · Aug 5, 2024 · 448ecbd · 448ecbd
1 parent 712405e
commit 448ecbd
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs
@@ -624,6 +624,10 @@ impl Scanner {
  extra_columns.push(ArrowField::new(DIST_COL, DataType::Float32, true));
  };
 
+ if self.full_text_query.is_some() {
+ extra_columns.push(ArrowField::new(SCORE_COL, DataType::Float32, true));
+ }
+
  if self.with_row_id || in_projection {
  extra_columns.push(ROW_ID_FIELD.clone());
  }
@@ -666,6 +670,11 @@ impl Scanner {
  output_expr.push((vector_expr, DIST_COL.to_string()));
  }
 
+ if self.full_text_query.is_some() && output_expr.iter().all(|(_, name)| name != SCORE_COL) {
+ let score_expr = expressions::col(SCORE_COL, &physical_schema)?;
+ output_expr.push((score_expr, SCORE_COL.to_string()));
+ }
+
  if self.with_row_id && output_expr.iter().all(|(_, name)| name != ROW_ID) {
  let row_id_expr = expressions::col(ROW_ID, &physical_schema)?;
  output_expr.push((row_id_expr, ROW_ID.to_string()));