Support sparse vector metadata value type

Author: Sicheng-Pan

chromadb/api/types.py

@@ -27,6 +27,7 @@
     Where,
     WhereDocumentOperator,
     WhereDocument,
+    SparseVector,
 )
 from inspect import signature
 from tenacity import retry
@@ -45,6 +46,9 @@
     "UpdateMetadata",
     "SearchRecord",
     "SearchResult",
+    "SparseVector",
+    "is_valid_sparse_vector",
+    "validate_sparse_vector",
 ]
 META_KEY_CHROMA_DOCUMENT = "chroma:document"
 T = TypeVar("T")
@@ -744,8 +748,69 @@ def validate_ids(ids: IDs) -> IDs:
     return ids
 
 
+def is_valid_sparse_vector(value: Any) -> bool:
+    """Check if a value looks like a SparseVector (has indices and values keys)."""
+    return isinstance(value, dict) and "indices" in value and "values" in value
+
+
+def validate_sparse_vector(value: Any) -> None:
+    """Validate that a value is a properly formed SparseVector.
+    
+    Args:
+        value: The value to validate as a SparseVector
+        
+    Raises:
+        ValueError: If the value is not a valid SparseVector
+    """
+    if not isinstance(value, dict):
+        raise ValueError(f"Expected SparseVector to be a dict, got {type(value).__name__}")
+    
+    if "indices" not in value or "values" not in value:
+        raise ValueError("SparseVector must have 'indices' and 'values' keys")
+    
+    indices = value.get("indices")
+    values = value.get("values")
+    
+    # Validate indices
+    if not isinstance(indices, list):
+        raise ValueError(
+            f"Expected SparseVector indices to be a list, got {type(indices).__name__}"
+        )
+    
+    # Validate values
+    if not isinstance(values, list):
+        raise ValueError(
+            f"Expected SparseVector values to be a list, got {type(values).__name__}"
+        )
+    
+    # Check lengths match
+    if len(indices) != len(values):
+        raise ValueError(
+            f"SparseVector indices and values must have the same length, "
+            f"got {len(indices)} indices and {len(values)} values"
+        )
+    
+    # Validate each index
+    for i, idx in enumerate(indices):
+        if not isinstance(idx, int):
+            raise ValueError(
+                f"SparseVector indices must be integers, got {type(idx).__name__} at position {i}"
+            )
+        if idx < 0:
+            raise ValueError(
+                f"SparseVector indices must be non-negative, got {idx} at position {i}"
+            )
+    
+    # Validate each value
+    for i, val in enumerate(values):
+        if not isinstance(val, (int, float)):
+            raise ValueError(
+                f"SparseVector values must be numbers, got {type(val).__name__} at position {i}"
+            )
+
+
 def validate_metadata(metadata: Metadata) -> Metadata:
-    """Validates metadata to ensure it is a dictionary of strings to strings, ints, floats or bools"""
+    """Validates metadata to ensure it is a dictionary of strings to strings, ints, floats, bools, or SparseVectors"""
     if not isinstance(metadata, dict) and metadata is not None:
         raise ValueError(
             f"Expected metadata to be a dict or None, got {type(metadata).__name__} as metadata"
@@ -765,18 +830,24 @@ def validate_metadata(metadata: Metadata) -> Metadata:
             raise TypeError(
                 f"Expected metadata key to be a str, got {key} which is a {type(key).__name__}"
             )
+        # Check if value is a SparseVector
+        if is_valid_sparse_vector(value):
+            try:
+                validate_sparse_vector(value)
+            except ValueError as e:
+                raise ValueError(f"Invalid SparseVector for key '{key}': {e}")
         # isinstance(True, int) evaluates to True, so we need to check for bools separately
-        if not isinstance(value, bool) and not isinstance(
+        elif not isinstance(value, bool) and not isinstance(
             value, (str, int, float, type(None))
         ):
             raise ValueError(
-                f"Expected metadata value to be a str, int, float, bool, or None, got {value} which is a {type(value).__name__}"
+                f"Expected metadata value to be a str, int, float, bool, SparseVector, or None, got {value} which is a {type(value).__name__}"
             )
     return metadata
 
 
 def validate_update_metadata(metadata: UpdateMetadata) -> UpdateMetadata:
-    """Validates metadata to ensure it is a dictionary of strings to strings, ints, floats or bools"""
+    """Validates metadata to ensure it is a dictionary of strings to strings, ints, floats, bools, or SparseVectors"""
     if not isinstance(metadata, dict) and metadata is not None:
         raise ValueError(
             f"Expected metadata to be a dict or None, got {type(metadata)}"
@@ -788,12 +859,18 @@ def validate_update_metadata(metadata: UpdateMetadata) -> UpdateMetadata:
     for key, value in metadata.items():
         if not isinstance(key, str):
             raise ValueError(f"Expected metadata key to be a str, got {key}")
+        # Check if value is a SparseVector
+        if is_valid_sparse_vector(value):
+            try:
+                validate_sparse_vector(value)
+            except ValueError as e:
+                raise ValueError(f"Invalid SparseVector for key '{key}': {e}")
         # isinstance(True, int) evaluates to True, so we need to check for bools separately
-        if not isinstance(value, bool) and not isinstance(
+        elif not isinstance(value, bool) and not isinstance(
             value, (str, int, float, type(None))
         ):
             raise ValueError(
-                f"Expected metadata value to be a str, int, or float, got {value}"
+                f"Expected metadata value to be a str, int, float, bool, SparseVector, or None, got {value}"
             )
     return metadata
 

chromadb/base_types.py

@@ -1,10 +1,22 @@
 from typing import Dict, List, Mapping, Optional, Sequence, Union
-from typing_extensions import Literal
+from typing_extensions import Literal, TypedDict
 import numpy as np
 from numpy.typing import NDArray
 
-Metadata = Mapping[str, Optional[Union[str, int, float, bool]]]
-UpdateMetadata = Mapping[str, Union[int, float, str, bool, None]]
+
+class SparseVector(TypedDict):
+    """Represents a sparse vector using parallel arrays for indices and values.
+    
+    Attributes:
+        indices: List of dimension indices (must be non-negative integers)
+        values: List of values corresponding to each index
+    """
+    indices: List[int]
+    values: List[float]
+
+
+Metadata = Mapping[str, Optional[Union[str, int, float, bool, SparseVector]]]
+UpdateMetadata = Mapping[str, Union[int, float, str, bool, SparseVector, None]]
 PyVector = Union[Sequence[float], Sequence[int]]
 Vector = NDArray[Union[np.int32, np.float32]]  # TODO: Specify that the vector is 1D
 # Metadata Query Grammar

chromadb/execution/expression/operator.py

@@ -1,7 +1,7 @@
 from dataclasses import dataclass, field
 from typing import Optional, List, Dict, Set, Any
 
-from chromadb.api.types import Embeddings, IDs, Include
+from chromadb.api.types import Embeddings, IDs, Include, SparseVector
 from chromadb.types import (
     Collection,
     RequestVersionContext,
@@ -232,11 +232,12 @@ def to_dict(self) -> Dict[str, Any]:
 @dataclass
 class SparseKnn(Rank):
     """Sparse KNN ranking"""
-    embedding: Dict[int, float]  # Sparse vector: index -> value
+    embedding: SparseVector  # Sparse vector with indices and values
     key: str  # No default for sparse KNN
     limit: int = 1024
 
     def to_dict(self) -> Dict[str, Any]:
+        # Convert SparseVector to the format expected by Rust API
         result = {"embedding": self.embedding, "key": self.key}
         if self.limit != 1024:
             result["limit"] = self.limit # type: ignore[assignment]