Rename to search

Author: Sicheng-Pan

idl/chromadb/proto/query_executor.proto

@@ -154,26 +154,26 @@ message ProjectOperator {
     repeated string fields = 1;
 }
 
-message RetrievePayload {
+message SearchPayload {
     FilterOperator filter = 1;
     ScoreOperator score = 2;
     LimitOperator limit = 3;
     ProjectOperator project = 4;
 }
 
-message RetrievePlan {
+message SearchPlan {
     ScanOperator scan = 1;
-    repeated RetrievePayload payloads = 2;
+    repeated SearchPayload payloads = 2;
 }
 
-message RetrieveResult {
+message SearchResult {
     repeated string results = 1;
 }
 
 service QueryExecutor {
     rpc Count(CountPlan) returns (CountResult) {}
     rpc Get(GetPlan) returns (GetResult) {}
     rpc KNN(KNNPlan) returns (KNNBatchResult) {}
-    rpc Retrieve(RetrievePlan) returns (RetrieveResult) {}
+    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, RetrieveResult},
-    plan::{Count, Get, Knn, Retrieve},
+    operator::{CountResult, GetResult, KnnBatchResult, SearchResult},
+    plan::{Count, Get, Knn, Search},
     ExecutorError,
 };
 
@@ -210,7 +210,7 @@ impl DistributedExecutor {
         Ok(res.into_inner().try_into()?)
     }
 
-    pub async fn retrieve(&mut self, plan: Retrieve) -> Result<RetrieveResult, ExecutorError> {
+    pub async fn search(&mut self, plan: Search) -> Result<SearchResult, ExecutorError> {
         // Get the collection ID from the plan
         let collection_id = &plan
             .scan
@@ -225,7 +225,7 @@ impl DistributedExecutor {
             .map_err(|e| ExecutorError::Internal(e.boxed()))?;
 
         // Convert plan to proto
-        let request: chroma_types::chroma_proto::RetrievePlan = plan.try_into()?;
+        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();
@@ -236,7 +236,7 @@ impl DistributedExecutor {
                     attempt_count.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
                 let is_retry = current_attempt > 0;
                 choose_query_client_weighted(&clients, &config, is_retry)?
-                    .retrieve(Request::new(request.clone()))
+                    .search(Request::new(request.clone()))
                     .await
             })
             .retry(self.backoff)

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, RetrieveResult,
+        Projection, ProjectionRecord, RecordDistance, SearchResult,
     },
-    plan::{Count, Get, Knn, Retrieve},
+    plan::{Count, Get, Knn, Search},
     CollectionAndSegments, CollectionUuid, ExecutorError, HnswSpace, SegmentType,
 };
 use std::{
@@ -306,9 +306,9 @@ impl LocalExecutor {
         })
     }
 
-    pub async fn retrieve(&mut self, _plan: Retrieve) -> Result<RetrieveResult, ExecutorError> {
+    pub async fn search(&mut self, _plan: Search) -> Result<SearchResult, ExecutorError> {
         Err(ExecutorError::NotImplemented(
-            "Retrieve operation is not implemented for local executor".to_string(),
+            "Search operation is not implemented for local executor".to_string(),
         ))
     }
 

rust/frontend/src/executor/mod.rs

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

rust/frontend/src/impls/service_based_frontend.rs

@@ -18,7 +18,7 @@ use chroma_sysdb::{GetCollectionsOptions, SysDb};
 use chroma_system::System;
 use chroma_types::{
     operator::{Filter, KnnBatch, KnnProjection, Limit, Projection, Scan},
-    plan::{Count, Get, Knn, Retrieve},
+    plan::{Count, Get, Knn, Search},
     AddCollectionRecordsError, AddCollectionRecordsRequest, AddCollectionRecordsResponse,
     Collection, CollectionUuid, CountCollectionsError, CountCollectionsRequest,
     CountCollectionsResponse, CountRequest, CountResponse, CreateCollectionError,
@@ -34,7 +34,7 @@ use chroma_types::{
     HeartbeatError, HeartbeatResponse, Include, KnnIndex, ListCollectionsRequest,
     ListCollectionsResponse, ListDatabasesError, ListDatabasesRequest, ListDatabasesResponse,
     Operation, OperationRecord, QueryError, QueryRequest, QueryResponse, ResetError, ResetResponse,
-    RetrieveRequest, RetrieveResponse, Segment, SegmentScope, SegmentType, SegmentUuid,
+    SearchRequest, SearchResponse, Segment, SegmentScope, SegmentType, SegmentUuid,
     UpdateCollectionError, UpdateCollectionRecordsError, UpdateCollectionRecordsRequest,
     UpdateCollectionRecordsResponse, UpdateCollectionRequest, UpdateCollectionResponse,
     UpdateTenantError, UpdateTenantRequest, UpdateTenantResponse, UpsertCollectionRecordsError,
@@ -63,7 +63,7 @@ struct Metrics {
     add_retries_counter: Counter<u64>,
     update_retries_counter: Counter<u64>,
     upsert_retries_counter: Counter<u64>,
-    retrieve_retries_counter: Counter<u64>,
+    search_retries_counter: Counter<u64>,
 }
 
 #[derive(Clone, Debug)]
@@ -103,7 +103,7 @@ impl ServiceBasedFrontend {
         let add_retries_counter = meter.u64_counter("add_retries").build();
         let update_retries_counter = meter.u64_counter("update_retries").build();
         let upsert_retries_counter = meter.u64_counter("upsert_retries").build();
-        let retrieve_retries_counter = meter.u64_counter("retrieve_retries").build();
+        let search_retries_counter = meter.u64_counter("search_retries").build();
         let metrics = Arc::new(Metrics {
             fork_retries_counter,
             delete_retries_counter,
@@ -113,7 +113,7 @@ impl ServiceBasedFrontend {
             add_retries_counter,
             update_retries_counter,
             upsert_retries_counter,
-            retrieve_retries_counter,
+            search_retries_counter,
         });
         // factor: 2.0,
         // min_delay_ms: 100,
@@ -1581,47 +1581,47 @@ impl ServiceBasedFrontend {
         res
     }
 
-    pub async fn retryable_retrieve(
+    pub async fn retryable_search(
         &mut self,
-        request: RetrieveRequest,
-    ) -> Result<RetrieveResponse, QueryError> {
+        request: SearchRequest,
+    ) -> Result<SearchResponse, QueryError> {
         // Get collection and segments once for all queries
         let collection_and_segments = self
             .collections_with_segments_provider
             .get_collection_with_segments(request.collection_id)
             .await
             .map_err(|err| QueryError::Other(Box::new(err) as Box<dyn ChromaError>))?;
 
-        // Create a single Retrieve plan with one scan and the payloads from the request
-        let retrieve_plan = Retrieve {
+        // Create a single Search plan with one scan and the payloads from the request
+        let search_plan = Search {
             scan: Scan {
                 collection_and_segments,
             },
-            payloads: request.retrievals,
+            payloads: request.searches,
         };
 
-        // Execute the single retrieve plan using the executor
-        let result = self.executor.retrieve(retrieve_plan).await?;
+        // Execute the single search plan using the executor
+        let result = self.executor.search(search_plan).await?;
 
-        Ok(RetrieveResponse {
+        Ok(SearchResponse {
             results: result.results,
         })
     }
 
-    pub async fn retrieve(
+    pub async fn search(
         &mut self,
-        request: RetrieveRequest,
-    ) -> Result<RetrieveResponse, QueryError> {
+        request: SearchRequest,
+    ) -> Result<SearchResponse, QueryError> {
         let retries = Arc::new(AtomicUsize::new(0));
-        let retrieve_to_retry = || {
+        let search_to_retry = || {
             let mut self_clone = self.clone();
             let request_clone = request.clone();
             let cache_clone = self
                 .collections_with_segments_provider
                 .collections_with_segments_cache
                 .clone();
             async move {
-                let res = self_clone.retryable_retrieve(request_clone).await;
+                let res = self_clone.retryable_search(request_clone).await;
                 match res {
                     Ok(res) => Ok(res),
                     Err(e) => {
@@ -1637,22 +1637,22 @@ impl ServiceBasedFrontend {
                 }
             }
         };
-        let res = retrieve_to_retry
+        let res = search_to_retry
             .retry(self.collections_with_segments_provider.get_retry_backoff())
             // NOTE: Transport level errors will manifest as unknown errors, and they should also be retried
             .when(|e| matches!(e.code(), ErrorCodes::NotFound | ErrorCodes::Unknown))
             .notify(|_, _| {
                 let retried = retries.fetch_add(1, Ordering::Relaxed);
                 if retried > 0 {
                     tracing::info!(
-                        "Retrying retrieve() request for collection {}",
+                        "Retrying search() request for collection {}",
                         request.collection_id
                     );
                 }
             })
             .await;
         self.metrics
-            .retrieve_retries_counter
+            .search_retries_counter
             .add(retries.load(Ordering::Relaxed) as u64, &[]);
         res
     }

rust/frontend/src/server.rs

@@ -10,7 +10,7 @@ use chroma_metering::{
     ExternalCollectionReadContext, MeteredFutureExt, ReadAction, StartRequest, WriteAction,
 };
 use chroma_system::System;
-use chroma_types::plan::RetrievePayload;
+use chroma_types::plan::SearchPayload;
 use chroma_types::{
     AddCollectionRecordsResponse, ChecklistResponse, Collection, CollectionConfiguration,
     CollectionMetadataUpdate, CollectionUuid, CountCollectionsRequest, CountCollectionsResponse,
@@ -22,7 +22,7 @@ use chroma_types::{
     HeartbeatResponse, IncludeList, InternalCollectionConfiguration,
     InternalUpdateCollectionConfiguration, ListCollectionsRequest, ListCollectionsResponse,
     ListDatabasesRequest, ListDatabasesResponse, Metadata, QueryRequest, QueryResponse,
-    RetrieveRequest, RetrieveResponse, UpdateCollectionConfiguration,
+    SearchRequest, SearchResponse, UpdateCollectionConfiguration,
     UpdateCollectionRecordsResponse, UpdateCollectionResponse, UpdateMetadata, UpdateTenantRequest,
     UpdateTenantResponse, UpsertCollectionRecordsResponse,
 };
@@ -125,7 +125,7 @@ pub struct Metrics {
     collection_count: Counter<u64>,
     collection_get: Counter<u64>,
     collection_query: Counter<u64>,
-    collection_retrieve: Counter<u64>,
+    collection_search: Counter<u64>,
 }
 
 impl Metrics {
@@ -159,7 +159,7 @@ impl Metrics {
             collection_count: meter.u64_counter("collection_count").build(),
             collection_get: meter.u64_counter("collection_get").build(),
             collection_query: meter.u64_counter("collection_query").build(),
-            collection_retrieve: meter.u64_counter("collection_retrieve").build(),
+            collection_search: meter.u64_counter("collection_search").build(),
         }
     }
 }
@@ -299,7 +299,7 @@ impl FrontendServer {
             )
             .route(
                 "/api/v2/tenants/{tenant}/databases/{database}/collections/{collection_id}/retrieve",
-                post(collection_retrieve),
+                post(collection_search),
             )
             .merge(docs_router)
             .with_state(self)
@@ -2175,41 +2175,41 @@ async fn collection_query(
 }
 
 #[derive(Debug, Clone, Deserialize, ToSchema)]
-pub struct RetrieveRequestPayload {
-    retrievals: Vec<RetrievePayload>,
+pub struct SearchRequestPayload {
+    searches: Vec<SearchPayload>,
 }
 
-/// Retrieve records from a collection with hybrid criterias.
+/// Search records from a collection with hybrid criterias.
 #[utoipa::path(
     post,
-    path = "/api/v2/tenants/{tenant}/databases/{database}/collections/{collection_id}/retrieve",
-    request_body = RetrieveRequestPayload,
+    path = "/api/v2/tenants/{tenant}/databases/{database}/collections/{collection_id}/search",
+    request_body = SearchRequestPayload,
     responses(
-        (status = 200, description = "Records retrieved from the collection", body = RetrieveResponse),
+        (status = 200, description = "Records searched from the collection", body = SearchResponse),
         (status = 401, description = "Unauthorized", body = ErrorResponse),
         (status = 404, description = "Collection not found", body = ErrorResponse),
         (status = 500, description = "Server error", body = ErrorResponse)
     ),
     params(
         ("tenant" = String, Path, description = "Tenant ID"),
         ("database" = String, Path, description = "Database name for the collection"),
-        ("collection_id" = String, Path, description = "Collection ID to retrieve records from")
+        ("collection_id" = String, Path, description = "Collection ID to search records from")
     )
 )]
-async fn collection_retrieve(
+async fn collection_search(
     headers: HeaderMap,
     Path((tenant, database, collection_id)): Path<(String, String, String)>,
     State(mut server): State<FrontendServer>,
-    TracedJson(payload): TracedJson<RetrieveRequestPayload>,
-) -> Result<Json<RetrieveResponse>, ServerError> {
-    server.metrics.collection_retrieve.add(
+    TracedJson(payload): TracedJson<SearchRequestPayload>,
+) -> Result<Json<SearchResponse>, ServerError> {
+    server.metrics.collection_search.add(
         1,
         &[
             KeyValue::new("tenant", tenant.clone()),
             KeyValue::new("collection_id", collection_id.clone()),
         ],
     );
-    // TODO: Maybe add AuthzAction:Retrieve
+    // TODO: Maybe add AuthzAction:Search
     let requester_identity = server
         .authenticate_and_authorize_collection(
             &headers,
@@ -2230,15 +2230,15 @@ async fn collection_retrieve(
         .map(|val| val.to_str().unwrap_or_default())
         .map(|val| val.to_string());
 
-    // TODO: Add quota enforcement for retrieve
+    // TODO: Add quota enforcement for search
     let _guard = server.scorecard_request(&[
         "op:read",
         format!("tenant:{}", tenant).as_str(),
         format!("collection:{}", collection_id).as_str(),
         format!("requester:{}", requester_identity.tenant).as_str(),
     ])?;
 
-    // TODO: Maybe add ReadAction::Retrieve
+    // TODO: Maybe add ReadAction::Search
     // Create a metering context
     let metering_context_container = if requester_identity.tenant == tenant {
         chroma_metering::create::<CollectionReadContext>(CollectionReadContext::new(
@@ -2265,14 +2265,14 @@ async fn collection_retrieve(
     });
 
     tracing::info!(
-        name: "collection_retrieve",
-        num_queries = payload.retrievals.len(),
+        name: "collection_search",
+        num_queries = payload.searches.len(),
     );
 
-    let request = RetrieveRequest::try_new(tenant, database, collection_id, payload.retrievals)?;
+    let request = SearchRequest::try_new(tenant, database, collection_id, payload.searches)?;
     let res = server
         .frontend
-        .retrieve(request)
+        .search(request)
         .meter(metering_context_container)
         .await?;
     Ok(Json(res))
@@ -2335,7 +2335,7 @@ impl Modify for ChromaTokenSecurityAddon {
         collection_count,
         collection_get,
         collection_query,
-        collection_retrieve,
+        collection_search,
     ),
     // Apply our new security scheme here
     modifiers(&ChromaTokenSecurityAddon)