[ENH] Use Idf if the index is BM25 (#5609)

## Description of changes

_Summarize the changes made by this PR._

- Improvements & Bug fixes
    - N/A
- New functionality
    - Add `bm25` flag in sparse vector config for schema
    - Scale query by IDF automatically if the key corresponds to BM25 index

## Test plan

_How are these changes tested?_

- [ ] Tests pass locally with `pytest` for python, `yarn test` for js, `cargo test` for rust

## Migration plan

_Are there any migrations, or any forwards/backwards compatibility changes needed in order to make sure this change deploys reliably?_

## Observability plan

_What is the plan to instrument and monitor this change?_

## Documentation Changes

_Are all docstrings for user-facing APIs updated if required? Do we need to make documentation changes in the_ [_docs section](https://github.com/chroma-core/chroma/tree/main/docs/docs.trychroma.com)?_

Author: Sicheng-Pan

clients/new-js/packages/chromadb/src/api/types.gen.ts

@@ -380,6 +380,10 @@ export type SparseVector = {
 };
 
 export type SparseVectorIndexConfig = {
+    /**
+     * Whether this embedding is BM25
+     */
+    bm25?: boolean | null;
     embedding_function?: null | EmbeddingFunctionConfiguration;
     /**
      * Key to source the sparse vector from

rust/frontend/src/server.rs

@@ -1989,11 +1989,13 @@ async fn collection_get(
         payload.offset.unwrap_or(0),
         payload.include,
     )?;
-    let res = server
-        .frontend
-        .get(request)
-        .meter(metering_context_container)
-        .await?;
+    let res = Box::pin(
+        server
+            .frontend
+            .get(request)
+            .meter(metering_context_container),
+    )
+    .await?;
     Ok(Json(res))
 }
 

rust/frontend/tests/proptest_helpers/frontend_under_test.rs

@@ -110,10 +110,8 @@ impl StateMachineTest for FrontendUnderTest {
                     // Update stats
                     {
                         if request.r#where.is_some() {
-                            let filtered_records = state
-                                .frontend
-                                .clone()
-                                .get(
+                            let filtered_records = Box::pin(
+                                state.frontend.clone().get(
                                     GetRequest::try_new(
                                         collection.tenant,
                                         collection.database,
@@ -125,9 +123,10 @@ impl StateMachineTest for FrontendUnderTest {
                                         IncludeList(vec![]),
                                     )
                                     .unwrap(),
-                                )
-                                .await
-                                .unwrap();
+                                ),
+                            )
+                            .await
+                            .unwrap();
 
                             STATS.with_borrow_mut(|stats| {
                                 stats.num_log_operations += filtered_records.ids.len()
@@ -199,7 +198,7 @@ impl StateMachineTest for FrontendUnderTest {
                         request.tenant_id = collection.tenant;
                         request.database_name = collection.database;
 
-                        state.frontend.get(request.clone()).await.unwrap()
+                        Box::pin(state.frontend.get(request.clone())).await.unwrap()
                     };
 
                     check_get_responses_are_close_to_equal(expected_result, received_result);
@@ -332,8 +331,8 @@ impl StateMachineTest for FrontendUnderTest {
                 )
                 .unwrap();
 
-            let received_results = frontend_under_test
-                .get(
+            let received_results = Box::pin(
+                frontend_under_test.get(
                     GetRequest::try_new(
                         collection_under_test.tenant,
                         collection_under_test.database,
@@ -345,9 +344,10 @@ impl StateMachineTest for FrontendUnderTest {
                         IncludeList::default_get(),
                     )
                     .unwrap(),
-                )
-                .await
-                .unwrap();
+                ),
+            )
+            .await
+            .unwrap();
 
             check_get_responses_are_close_to_equal(expected_results, received_results);
         });

rust/python_bindings/src/bindings.rs

@@ -623,7 +623,7 @@ impl Bindings {
         let mut frontend_clone = self.frontend.clone();
         let result = py.allow_threads(move || {
             self.runtime
-                .block_on(async { frontend_clone.get(request).await })
+                .block_on(async { Box::pin(frontend_clone.get(request)).await })
         })?;
         Ok(result)
     }

rust/types/src/collection_schema.rs

@@ -311,6 +311,7 @@ impl InternalSchema {
                     config: SparseVectorIndexConfig {
                         embedding_function: Some(EmbeddingFunctionConfiguration::Legacy),
                         source_key: None,
+                        bm25: Some(false),
                     },
                 }),
             }),
@@ -909,6 +910,7 @@ impl InternalSchema {
                 .clone()
                 .or(default.embedding_function.clone()),
             source_key: user.source_key.clone().or(default.source_key.clone()),
+            bm25: user.bm25.or(default.bm25),
         })
     }
 
@@ -1436,6 +1438,9 @@ pub struct SparseVectorIndexConfig {
     /// Key to source the sparse vector from
     #[serde(skip_serializing_if = "Option::is_none")]
     pub source_key: Option<String>,
+    /// Whether this embedding is BM25
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub bm25: Option<bool>,
 }
 
 #[derive(Clone, Debug, PartialEq, Serialize, Deserialize, ToSchema)]
@@ -1791,6 +1796,7 @@ mod tests {
                         config: SparseVectorIndexConfig {
                             embedding_function: Some(EmbeddingFunctionConfiguration::Legacy),
                             source_key: None,
+                            bm25: None,
                         },
                     }),
                 }),
@@ -2072,11 +2078,13 @@ mod tests {
         let default_config = SparseVectorIndexConfig {
             embedding_function: Some(EmbeddingFunctionConfiguration::Legacy),
             source_key: Some("default_sparse_key".to_string()),
+            bm25: None,
         };
 
         let user_config = SparseVectorIndexConfig {
             embedding_function: None,                        // Will use default
             source_key: Some("user_sparse_key".to_string()), // Override
+            bm25: None,
         };
 
         let result =

rust/worker/src/execution/orchestration/sparse_knn.rs

@@ -209,7 +209,29 @@ impl Orchestrator for SparseKnnOrchestrator {
         &mut self,
         ctx: &ComponentContext<Self>,
     ) -> Vec<(TaskMessage, Option<Span>)> {
-        if self.use_bm25 {
+        let use_bm25 = self.use_bm25
+            || self
+                .collection_and_segments
+                .collection
+                .schema
+                .as_ref()
+                .is_some_and(|schema| {
+                    if let Some(flag) = schema.key_overrides.get(&self.key).and_then(|uvt| {
+                        uvt.sparse_vector.as_ref().and_then(|vt| {
+                            vt.sparse_vector_index
+                                .as_ref()
+                                .and_then(|it| it.config.bm25)
+                        })
+                    }) {
+                        return flag;
+                    }
+                    schema.defaults.sparse_vector.as_ref().is_some_and(|vt| {
+                        vt.sparse_vector_index
+                            .as_ref()
+                            .is_some_and(|it| it.config.bm25.unwrap_or_default())
+                    })
+                });
+        if use_bm25 {
             let idf_task = wrap(
                 Box::new(Idf {
                     query: self.query.clone(),