[ENH] Transpose search api response (#5414)

## Description of changes

_Summarize the changes made by this PR._

- Improvements & Bug fixes
  - Transpose the search response from row major to column major for consistency with our existing endpoints
- New functionality
  - N/A

## 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

chromadb/api/async_fastapi.py

@@ -41,7 +41,6 @@
     GetResult,
     QueryResult,
     SearchResult,
-    SearchRecord,
     CollectionMetadata,
     validate_batch,
     convert_np_embeddings_to_list,
@@ -420,20 +419,15 @@ async def _search(
             json=payload,
         )
 
-        # Parse response into SearchResult
-        results = []
-        for batch_results in resp_json.get("results", []):
-            batch = []
-            for record in batch_results:
-                batch.append(SearchRecord(
-                    id=record["id"],
-                    document=record.get("document"),
-                    embedding=record.get("embedding"),
-                    metadata=record.get("metadata"),
-                    score=record.get("score"),
-                ))
-            results.append(batch)
-        return results
+        # Return the column-major format directly
+        return SearchResult(
+            ids=resp_json.get("ids", []),
+            documents=resp_json.get("documents", []),
+            embeddings=resp_json.get("embeddings", []),
+            metadatas=resp_json.get("metadatas", []),
+            scores=resp_json.get("scores", []),
+            select=resp_json.get("select", [])
+        )
 
     @trace_method("AsyncFastAPI.delete_collection", OpenTelemetryGranularity.OPERATION)
     @override

chromadb/api/fastapi.py

@@ -32,7 +32,6 @@
     GetResult,
     QueryResult,
     SearchResult,
-    SearchRecord,
     CollectionMetadata,
     validate_batch,
     convert_np_embeddings_to_list,
@@ -377,20 +376,15 @@ def _search(
             json=payload,
         )
 
-        # Parse response into SearchResult
-        results = []
-        for batch_results in resp_json.get("results", []):
-            batch = []
-            for record in batch_results:
-                batch.append(SearchRecord(
-                    id=record["id"],
-                    document=record.get("document"),
-                    embedding=record.get("embedding"),
-                    metadata=record.get("metadata"),
-                    score=record.get("score"),
-                ))
-            results.append(batch)
-        return results
+        # Return the column-major format directly
+        return SearchResult(
+            ids=resp_json.get("ids", []),
+            documents=resp_json.get("documents", []),
+            embeddings=resp_json.get("embeddings", []),
+            metadatas=resp_json.get("metadatas", []),
+            scores=resp_json.get("scores", []),
+            select=resp_json.get("select", [])
+        )
 
     @trace_method("FastAPI.delete_collection", OpenTelemetryGranularity.OPERATION)
     @override

chromadb/api/models/AsyncCollection.py

@@ -305,8 +305,13 @@ async def search(
                 - select: Select configuration for fields to return (defaults to empty)
         
         Returns:
-            SearchResult: List of search results for each search payload.
-                         Each result is a list of SearchRecord objects.
+            SearchResult: Column-major format response with:
+                - ids: List of result IDs for each search payload
+                - documents: Optional documents for each payload
+                - embeddings: Optional embeddings for each payload
+                - metadatas: Optional metadata for each payload
+                - scores: Optional scores for each payload
+                - select: List of selected fields for each payload
         
         Raises:
             NotImplementedError: For local/segment API implementations

chromadb/api/models/Collection.py

@@ -309,8 +309,13 @@ def search(
                 - select: Select configuration for fields to return (defaults to empty)
         
         Returns:
-            SearchResult: List of search results for each search payload.
-                         Each result is a list of SearchRecord objects.
+            SearchResult: Column-major format response with:
+                - ids: List of result IDs for each search payload
+                - documents: Optional documents for each payload
+                - embeddings: Optional embeddings for each payload
+                - metadatas: Optional metadata for each payload
+                - scores: Optional scores for each payload
+                - select: List of selected fields for each payload
         
         Raises:
             NotImplementedError: For local/segment API implementations

chromadb/api/types.py

@@ -10,6 +10,7 @@
     cast,
     Literal,
     get_args,
+    TYPE_CHECKING,
 )
 from numpy.typing import NDArray
 import numpy as np
@@ -29,6 +30,9 @@
     WhereDocument,
     SparseVector,
 )
+
+if TYPE_CHECKING:
+    from chromadb.execution.expression.operator import SelectField
 from inspect import signature
 from tenacity import retry
 from abc import abstractmethod
@@ -44,7 +48,6 @@
     "WhereDocument",
     "UpdateCollectionMetadata",
     "UpdateMetadata",
-    "SearchRecord",
     "SearchResult",
     "SparseVector",
     "is_valid_sparse_vector",
@@ -490,17 +493,26 @@ class QueryResult(TypedDict):
     included: Include
 
 
-class SearchRecord(TypedDict):
-    """Individual record returned from a search operation"""
-    id: str
-    document: Optional[str]
-    embedding: Optional[List[float]]
-    metadata: Optional[Metadata]
-    score: Optional[float]
-
-
-# SearchResult is a list of results for each search payload
-SearchResult = List[List[SearchRecord]]
+class SearchResult(TypedDict):
+    """Column-major response from the search API matching Rust SearchResponse structure.
+    
+    This is the format returned to users:
+    - ids: Always present, list of result IDs for each search payload
+    - documents: Optional per payload, None if not requested
+    - embeddings: Optional per payload, None if not requested
+    - metadatas: Optional per payload, None if not requested
+    - scores: Optional per payload, None if not requested
+    - select: List of selected fields for each payload (sorted)
+    
+    Each top-level list index corresponds to a search payload.
+    Within each payload, the inner lists are aligned by record index.
+    """
+    ids: List[List[str]]
+    documents: List[Optional[List[Optional[str]]]]
+    embeddings: List[Optional[List[Optional[List[float]]]]]
+    metadatas: List[Optional[List[Optional[Dict[str, Any]]]]]
+    scores: List[Optional[List[Optional[float]]]]
+    select: List[List[Union["SelectField", str]]]  # List of SelectField enums or string field names for each payload
 
 
 class UpdateRequest(TypedDict):

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

@@ -205,20 +205,21 @@ export type SearchPayload = {
     };
 };
 
-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>>;
+    documents: Array<Array<string | null> | null>;
+    embeddings: Array<Array<Array<number> | null> | null>;
+    ids: Array<Array<string>>;
+    metadatas: Array<Array<null | HashMap> | null>;
+    scores: Array<Array<number | null> | null>;
+    select: Array<Array<SelectField>>;
+};
+
+export type SelectField = 'Document' | 'Embedding' | 'Metadata' | 'Score' | {
+    MetadataField: string;
 };
 
 export type SpannConfiguration = {

rust/frontend/src/impls/service_based_frontend.rs

@@ -1641,6 +1641,8 @@ impl ServiceBasedFrontend {
         }
 
         // Create a single Search plan with one scan and the payloads from the request
+        // Clone the searches to use them later for aggregating select fields
+        let searches_for_select = request.searches.clone();
         let search_plan = Search {
             scan: Scan {
                 collection_and_segments,
@@ -1691,13 +1693,7 @@ impl ServiceBasedFrontend {
             },
         }
 
-        Ok(SearchResponse {
-            results: result
-                .results
-                .into_iter()
-                .map(|result| result.records)
-                .collect(),
-        })
+        Ok((result, searches_for_select).into())
     }
 
     pub async fn search(&mut self, request: SearchRequest) -> Result<SearchResponse, QueryError> {

rust/types/src/api_types.rs

@@ -5,7 +5,8 @@ use crate::operator::GetResult;
 use crate::operator::KnnBatchResult;
 use crate::operator::KnnProjectionRecord;
 use crate::operator::ProjectionRecord;
-use crate::operator::SearchRecord;
+use crate::operator::SearchResult;
+use crate::operator::SelectField;
 use crate::plan::PlanToProtoError;
 use crate::plan::SearchPayload;
 use crate::validators::{
@@ -1869,7 +1870,84 @@ impl SearchRequest {
 
 #[derive(Clone, Deserialize, Serialize, ToSchema, Debug)]
 pub struct SearchResponse {
-    pub results: Vec<Vec<SearchRecord>>,
+    pub ids: Vec<Vec<String>>,
+    pub documents: Vec<Option<Vec<Option<String>>>>,
+    pub embeddings: Vec<Option<Vec<Option<Vec<f32>>>>>,
+    pub metadatas: Vec<Option<Vec<Option<Metadata>>>>,
+    pub scores: Vec<Option<Vec<Option<f32>>>>,
+    pub select: Vec<Vec<SelectField>>,
+}
+
+impl From<(SearchResult, Vec<SearchPayload>)> for SearchResponse {
+    fn from((result, payloads): (SearchResult, Vec<SearchPayload>)) -> Self {
+        let num_payloads = payloads.len();
+        let mut res = Self {
+            ids: Vec::with_capacity(num_payloads),
+            documents: Vec::with_capacity(num_payloads),
+            embeddings: Vec::with_capacity(num_payloads),
+            metadatas: Vec::with_capacity(num_payloads),
+            scores: Vec::with_capacity(num_payloads),
+            select: Vec::with_capacity(num_payloads),
+        };
+
+        for (payload_result, payload) in result.results.into_iter().zip(payloads) {
+            // Get the sorted select fields for this payload
+            let mut payload_select = Vec::from_iter(payload.select.fields.iter().cloned());
+            payload_select.sort();
+
+            let num_records = payload_result.records.len();
+            let mut ids = Vec::with_capacity(num_records);
+            let mut documents = Vec::with_capacity(num_records);
+            let mut embeddings = Vec::with_capacity(num_records);
+            let mut metadatas = Vec::with_capacity(num_records);
+            let mut scores = Vec::with_capacity(num_records);
+
+            for record in payload_result.records {
+                ids.push(record.id);
+                documents.push(record.document);
+                embeddings.push(record.embedding);
+                metadatas.push(record.metadata);
+                scores.push(record.score);
+            }
+
+            res.ids.push(ids);
+            res.select.push(payload_select.clone());
+
+            // Push documents if requested by this payload, otherwise None
+            res.documents.push(
+                payload_select
+                    .binary_search(&SelectField::Document)
+                    .is_ok()
+                    .then_some(documents),
+            );
+
+            // Push embeddings if requested by this payload, otherwise None
+            res.embeddings.push(
+                payload_select
+                    .binary_search(&SelectField::Embedding)
+                    .is_ok()
+                    .then_some(embeddings),
+            );
+
+            // Push metadatas if requested by this payload, otherwise None
+            // Include if either SelectField::Metadata is present or any SelectField::MetadataField(_)
+            let has_metadata = payload_select.binary_search(&SelectField::Metadata).is_ok()
+                || payload_select
+                    .last()
+                    .is_some_and(|field| matches!(field, SelectField::MetadataField(_)));
+            res.metadatas.push(has_metadata.then_some(metadatas));
+
+            // Push scores if requested by this payload, otherwise None
+            res.scores.push(
+                payload_select
+                    .binary_search(&SelectField::Score)
+                    .is_ok()
+                    .then_some(scores),
+            );
+        }
+
+        res
+    }
 }
 
 #[derive(Error, Debug)]

rust/types/src/execution/operator.rs

@@ -942,7 +942,7 @@ impl TryFrom<Rank> for chroma_proto::Rank {
     }
 }
 
-#[derive(Clone, Debug, PartialEq, Eq, Hash)]
+#[derive(Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ToSchema)]
 pub enum SelectField {
     // Predefined fields
     Document,