[ENH] Implement new search endpoint (#5323)

## Description of changes

_Summarize the changes made by this PR._

- Improvements & Bug fixes
  - N/A
- New functionality
  - Introduces a new `search` endpoint for hybrid search

## 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/scripts/gen-api.ts

@@ -36,8 +36,8 @@ const main = async () => {
 
     // Fix the HashMap type to include null and remove duplicate number
     typesContent = typesContent.replace(
-      /export type HashMap = \{\s*\[key: string\]: boolean \| number \| number \| string;\s*\};/,
-      "export type HashMap = {\n  [key: string]: boolean | number | string | null;\n};",
+      /export type HashMap = \{\s*\[key: string\]: boolean \| number \| number \| string \| SparseVector;\s*};/,
+      "export type HashMap = {\n  [key: string]: boolean | number | string | SparseVector | null;\n};",
     );
 
     await writeFile(typesPath, typesContent);

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

@@ -1,7 +1,7 @@
 // This file is auto-generated by @hey-api/openapi-ts
 
 import type { Options as ClientOptions, TDataShape, Client } from '@hey-api/client-fetch';
-import type { GetUserIdentityData, GetUserIdentityResponse2, GetUserIdentityError, GetCollectionByCrnData, GetCollectionByCrnResponse, GetCollectionByCrnError, HealthcheckData, HealthcheckResponse, HealthcheckError, HeartbeatData, HeartbeatResponse2, HeartbeatError, PreFlightChecksData, PreFlightChecksResponse, PreFlightChecksError, ResetData, ResetResponse, ResetError, CreateTenantData, CreateTenantResponse2, CreateTenantError, GetTenantData, GetTenantResponse2, GetTenantError, UpdateTenantData, UpdateTenantResponse2, UpdateTenantError, ListDatabasesData, ListDatabasesResponse, ListDatabasesError, CreateDatabaseData, CreateDatabaseResponse2, CreateDatabaseError, DeleteDatabaseData, DeleteDatabaseResponse2, DeleteDatabaseError, GetDatabaseData, GetDatabaseResponse, GetDatabaseError, ListCollectionsData, ListCollectionsResponse, ListCollectionsError, CreateCollectionData, CreateCollectionResponse, CreateCollectionError, DeleteCollectionData, DeleteCollectionResponse, DeleteCollectionError, GetCollectionData, GetCollectionResponse, GetCollectionError, UpdateCollectionData, UpdateCollectionResponse2, UpdateCollectionError, CollectionAddData, CollectionAddResponse, CollectionCountData, CollectionCountResponse, CollectionCountError, CollectionDeleteData, CollectionDeleteResponse, CollectionDeleteError, ForkCollectionData, ForkCollectionResponse, ForkCollectionError, CollectionGetData, CollectionGetResponse, CollectionGetError, CollectionQueryData, CollectionQueryResponse, CollectionQueryError, CollectionUpdateData, CollectionUpdateResponse, CollectionUpsertData, CollectionUpsertResponse, CollectionUpsertError, CountCollectionsData, CountCollectionsResponse, CountCollectionsError, VersionData, VersionResponse } from './types.gen';
+import type { GetUserIdentityData, GetUserIdentityResponse2, GetUserIdentityError, GetCollectionByCrnData, GetCollectionByCrnResponse, GetCollectionByCrnError, HealthcheckData, HealthcheckResponse, HealthcheckError, HeartbeatData, HeartbeatResponse2, HeartbeatError, PreFlightChecksData, PreFlightChecksResponse, PreFlightChecksError, ResetData, ResetResponse, ResetError, CreateTenantData, CreateTenantResponse2, CreateTenantError, GetTenantData, GetTenantResponse2, GetTenantError, UpdateTenantData, UpdateTenantResponse2, UpdateTenantError, ListDatabasesData, ListDatabasesResponse, ListDatabasesError, CreateDatabaseData, CreateDatabaseResponse2, CreateDatabaseError, DeleteDatabaseData, DeleteDatabaseResponse2, DeleteDatabaseError, GetDatabaseData, GetDatabaseResponse, GetDatabaseError, ListCollectionsData, ListCollectionsResponse, ListCollectionsError, CreateCollectionData, CreateCollectionResponse, CreateCollectionError, DeleteCollectionData, DeleteCollectionResponse, DeleteCollectionError, GetCollectionData, GetCollectionResponse, GetCollectionError, UpdateCollectionData, UpdateCollectionResponse2, UpdateCollectionError, CollectionAddData, CollectionAddResponse, CollectionCountData, CollectionCountResponse, CollectionCountError, CollectionDeleteData, CollectionDeleteResponse, CollectionDeleteError, ForkCollectionData, ForkCollectionResponse, ForkCollectionError, CollectionGetData, CollectionGetResponse, CollectionGetError, CollectionQueryData, CollectionQueryResponse, CollectionQueryError, CollectionSearchData, CollectionSearchResponse, CollectionSearchError, CollectionUpdateData, CollectionUpdateResponse, CollectionUpsertData, CollectionUpsertResponse, CollectionUpsertError, CountCollectionsData, CountCollectionsResponse, CountCollectionsError, VersionData, VersionResponse } from './types.gen';
 import { client as _heyApiClient } from './client.gen';
 
 export type Options<TData extends TDataShape = TDataShape, ThrowOnError extends boolean = boolean> = ClientOptions<TData, ThrowOnError> & {
@@ -299,6 +299,20 @@ export class DefaultService {
         });
     }
 
+    /**
+     * Search records from a collection with hybrid criterias.
+     */
+    public static collectionSearch<ThrowOnError extends boolean = true>(options: Options<CollectionSearchData, ThrowOnError>) {
+        return (options.client ?? _heyApiClient).post<CollectionSearchResponse, CollectionSearchError, ThrowOnError>({
+            url: '/api/v2/tenants/{tenant}/databases/{database}/collections/{collection_id}/search',
+            ...options,
+            headers: {
+                'Content-Type': 'application/json',
+                ...options?.headers
+            }
+        });
+    }
+    
     /**
      * Updates records in a collection by ID.
      */

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

@@ -131,7 +131,7 @@ export type GetUserIdentityResponse = {
 };
 
 export type HashMap = {
-    [key: string]: boolean | number | number | string | SparseVector;
+  [key: string]: boolean | number | string | SparseVector | null;
 };
 
 export type HeartbeatResponse = {
@@ -175,6 +175,52 @@ export type RawWhereFields = {
     where_document?: unknown;
 };
 
+/**
+ * Payload for hybrid search
+ */
+export type SearchPayload = {
+    /**
+     * Filter criteria for search
+     */
+    filter: {
+        query_ids?: Array<string>;
+        where_clause?: {
+            [key: string]: unknown;
+        };
+    };
+    limit: {
+        fetch?: number;
+        skip: number;
+    };
+    /**
+     * Ranking expression for hybrid search
+     */
+    rank: {
+        [key: string]: {
+            [key: string]: unknown;
+        };
+    };
+    select: {
+        fields: Array<string>;
+    };
+};
+
+export type SearchRecord = {
+    document?: string | null;
+    embedding?: Array<number> | null;
+    id: string;
+    metadata?: null | HashMap;
+    score?: number | null;
+};
+
+export type SearchRequestPayload = {
+    searches: Array<SearchPayload>;
+};
+
+export type SearchResponse = {
+    results: Array<Array<SearchRecord>>;
+};
+
 export type SpannConfiguration = {
     ef_construction?: number | null;
     ef_search?: number | null;
@@ -1205,6 +1251,52 @@ export type CollectionQueryResponses = {
 
 export type CollectionQueryResponse = CollectionQueryResponses[keyof CollectionQueryResponses];
 
+export type CollectionSearchData = {
+    body: SearchRequestPayload;
+    path: {
+        /**
+         * Tenant ID
+         */
+        tenant: string;
+        /**
+         * Database name for the collection
+         */
+        database: string;
+        /**
+         * Collection ID to search records from
+         */
+        collection_id: string;
+    };
+    query?: never;
+    url: '/api/v2/tenants/{tenant}/databases/{database}/collections/{collection_id}/search';
+};
+
+export type CollectionSearchErrors = {
+    /**
+     * Unauthorized
+     */
+    401: ErrorResponse;
+    /**
+     * Collection not found
+     */
+    404: ErrorResponse;
+    /**
+     * Server error
+     */
+    500: ErrorResponse;
+};
+
+export type CollectionSearchError = CollectionSearchErrors[keyof CollectionSearchErrors];
+
+export type CollectionSearchResponses = {
+    /**
+     * Records searched from the collection
+     */
+    200: SearchResponse;
+};
+
+export type CollectionSearchResponse = CollectionSearchResponses[keyof CollectionSearchResponses];
+
 export type CollectionUpdateData = {
     body: UpdateCollectionRecordsPayload;
     path: {

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

@@ -1,4 +1,4 @@
-import { GetUserIdentityResponse, Include } from "./api";
+import { GetUserIdentityResponse, Include, SparseVector } from "./api";
 
 /**
  * User identity information including tenant and database access.
@@ -11,14 +11,17 @@ export type UserIdentity = GetUserIdentityResponse;
  */
 export type CollectionMetadata = Record<
   string,
-  boolean | number | string | null
+  boolean | number | string | SparseVector | null
 >;
 
 /**
  * Metadata that can be associated with individual records.
  * Values must be boolean, number, or string types.
  */
-export type Metadata = Record<string, boolean | number | string | null>;
+export type Metadata = Record<
+  string,
+  boolean | number | string | SparseVector | null
+>;
 
 /**
  * Base interface for record sets containing optional fields.

idl/chromadb/proto/query_executor.proto

@@ -89,9 +89,88 @@ message KNNBatchResult {
     uint64 pulled_log_bytes = 2;
 }
 
+message QueryVector {
+    oneof vector {
+        Vector dense = 1;
+        SparseVector sparse = 2;
+    }
+}
+
+message Rank {
+    message Division {
+        Rank left = 1;
+        Rank right = 2;
+    }
+
+    message Knn {
+        QueryVector embedding = 1;
+        string key = 2;
+        uint32 limit = 3;
+        optional float default = 4;
+        bool ordinal = 5;
+    }
+    
+    message Subtraction {
+        Rank left = 1;
+        Rank right = 2;
+    }
+
+    message RankList {
+        repeated Rank ranks = 1;
+    }
+
+    oneof rank {
+        Rank absolute = 1;
+        Division division = 2;
+        Rank exponentiation = 3;
+        Knn knn = 4;
+        Rank logarithm = 5;
+        RankList maximum = 6;
+        RankList minimum = 7;
+        RankList multiplication = 8;
+        Subtraction subtraction = 9;
+        RankList summation = 10;
+        float value = 11;
+    }
+}
+
+message SelectOperator {
+    repeated string fields = 1;
+}
+
+message SearchPayload {
+    FilterOperator filter = 1;
+    Rank rank = 2;
+    LimitOperator limit = 3;
+    SelectOperator select = 4;
+}
+
+message SearchPlan {
+    ScanOperator scan = 1;
+    repeated SearchPayload payloads = 2;
+}
+
+message SearchRecord {
+    string id = 1;
+    optional string document = 2;
+    optional Vector embedding = 3;
+    optional UpdateMetadata metadata = 4;
+    optional float score = 5;
+}
+
+message SearchPayloadResult {
+    repeated SearchRecord records = 1;
+}
+
+message SearchResult {
+    repeated SearchPayloadResult results = 1;
+    uint64 pulled_log_bytes = 2;
+}
+
 service QueryExecutor {
     rpc Count(CountPlan) returns (CountResult) {}
     rpc Get(GetPlan) returns (GetResult) {}
     rpc KNN(KNNPlan) returns (KNNBatchResult) {}
+    rpc Search(SearchPlan) returns (SearchResult) {}
 }
 

rust/frontend/src/executor/distributed.rs

@@ -15,8 +15,8 @@ use chroma_system::System;
 use chroma_types::chroma_proto::query_executor_client::QueryExecutorClient;
 use chroma_types::SegmentType;
 use chroma_types::{
-    operator::{CountResult, GetResult, KnnBatchResult},
-    plan::{Count, Get, Knn},
+    operator::{CountResult, GetResult, KnnBatchResult, SearchResult},
+    plan::{Count, Get, Knn, Search},
     ExecutorError,
 };
 
@@ -210,6 +210,42 @@ impl DistributedExecutor {
         Ok(res.into_inner().try_into()?)
     }
 
+    pub async fn search(&mut self, plan: Search) -> Result<SearchResult, ExecutorError> {
+        // Get the collection ID from the plan
+        let collection_id = &plan
+            .scan
+            .collection_and_segments
+            .collection
+            .collection_id
+            .to_string();
+
+        let clients = self
+            .client_assigner
+            .clients(collection_id)
+            .map_err(|e| ExecutorError::Internal(e.boxed()))?;
+
+        // Convert plan to proto
+        let request: chroma_types::chroma_proto::SearchPlan = plan.try_into()?;
+
+        let attempt_count = std::sync::Arc::new(std::sync::atomic::AtomicU32::new(0));
+        let config = self.client_selection_config.clone();
+        let res = {
+            let attempt_count = attempt_count.clone();
+            (|| async {
+                let current_attempt =
+                    attempt_count.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
+                let is_retry = current_attempt > 0;
+                choose_query_client_weighted(&clients, &config, is_retry)?
+                    .search(Request::new(request.clone()))
+                    .await
+            })
+            .retry(self.backoff)
+            .when(is_retryable_error)
+            .await?
+        };
+        Ok(res.into_inner().try_into()?)
+    }
+
     pub async fn is_ready(&self) -> bool {
         !self.client_assigner.is_empty()
     }

rust/frontend/src/executor/local.rs

@@ -12,9 +12,9 @@ use chroma_system::ComponentHandle;
 use chroma_types::{
     operator::{
         CountResult, Filter, GetResult, KnnBatchResult, KnnProjectionOutput, KnnProjectionRecord,
-        Projection, ProjectionRecord, RecordDistance,
+        Limit, Projection, ProjectionRecord, RecordDistance, SearchResult,
     },
-    plan::{Count, Get, Knn},
+    plan::{Count, Get, Knn, Search},
     CollectionAndSegments, CollectionUuid, ExecutorError, HnswSpace, SegmentType,
 };
 use std::{
@@ -161,7 +161,10 @@ impl LocalExecutor {
                 let filter_plan = Get {
                     scan: plan.scan.clone(),
                     filter: filter.clone(),
-                    limit: Default::default(),
+                    limit: Limit {
+                        skip: 0,
+                        fetch: None,
+                    },
                     proj: Default::default(),
                 };
 
@@ -264,7 +267,10 @@ impl LocalExecutor {
                     query_ids: Some(returned_user_ids),
                     where_clause: None,
                 },
-                limit: Default::default(),
+                limit: Limit {
+                    skip: 0,
+                    fetch: None,
+                },
                 proj: Projection {
                     document: plan.proj.projection.document,
                     embedding: false,
@@ -306,6 +312,12 @@ impl LocalExecutor {
         })
     }
 
+    pub async fn search(&mut self, _plan: Search) -> Result<SearchResult, ExecutorError> {
+        Err(ExecutorError::NotImplemented(
+            "Search operation is not implemented for local executor".to_string(),
+        ))
+    }
+
     pub async fn reset(&mut self) -> Result<(), Box<dyn ChromaError>> {
         self.hnsw_manager.reset().await.map_err(|err| err.boxed())?;
         Ok(())

rust/frontend/src/executor/mod.rs

@@ -1,6 +1,6 @@
 use chroma_types::{
-    operator::{CountResult, GetResult, KnnBatchResult},
-    plan::{Count, Get, Knn},
+    operator::{CountResult, GetResult, KnnBatchResult, SearchResult},
+    plan::{Count, Get, Knn, Search},
     ExecutorError, SegmentType,
 };
 use distributed::DistributedExecutor;
@@ -38,6 +38,12 @@ impl Executor {
             Executor::Local(local_executor) => local_executor.knn(plan).await,
         }
     }
+    pub async fn search(&mut self, plan: Search) -> Result<SearchResult, ExecutorError> {
+        match self {
+            Executor::Distributed(distributed_executor) => distributed_executor.search(plan).await,
+            Executor::Local(local_executor) => local_executor.search(plan).await,
+        }
+    }
     pub async fn is_ready(&self) -> bool {
         match self {
             Executor::Distributed(distributed_executor) => distributed_executor.is_ready().await,