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,10 @@ 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
-            }
-            Executor::Local(local_executor) => local_executor.retrieve(plan).await,
+            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 {

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,44 @@ 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(
-        &mut self,
-        request: RetrieveRequest,
-    ) -> Result<RetrieveResponse, QueryError> {
+    pub async fn search(&mut self, 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 +1634,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
     }