Refactor message parsing for more actions

Author: jhamon

pinecone/grpc/utils.py

@@ -1,4 +1,6 @@
-from typing import Any
+from __future__ import annotations
+
+from typing import Any, TYPE_CHECKING
 from google.protobuf import json_format
 from google.protobuf.message import Message
 
@@ -28,6 +30,17 @@
 
 from google.protobuf.struct_pb2 import Struct
 
+if TYPE_CHECKING:
+    from pinecone.core.grpc.protos.db_data_2025_10_pb2 import (
+        FetchResponse as ProtoFetchResponse,
+        FetchByMetadataResponse as ProtoFetchByMetadataResponse,
+        QueryResponse as ProtoQueryResponse,
+        UpsertResponse as ProtoUpsertResponse,
+        UpdateResponse as ProtoUpdateResponse,
+        NamespaceDescription as ProtoNamespaceDescription,
+        ListNamespacesResponse as ProtoListNamespacesResponse,
+    )
+
 
 def _generate_request_id() -> str:
     return str(uuid.uuid4())
@@ -53,7 +66,7 @@ def parse_sparse_values(sparse_values: dict | None) -> SparseValues:
 
 
 def parse_fetch_response(
-    response: Message, initial_metadata: dict[str, str] | None = None
+    response: "ProtoFetchResponse", initial_metadata: dict[str, str] | None = None
 ) -> FetchResponse:
     """Parse a FetchResponse protobuf message directly without MessageToDict conversion.
 
@@ -105,36 +118,61 @@ def parse_fetch_response(
 
 
 def parse_fetch_by_metadata_response(
-    response: Message, initial_metadata: dict[str, str] | None = None
+    response: "ProtoFetchByMetadataResponse", initial_metadata: dict[str, str] | None = None
 ) -> FetchByMetadataResponse:
-    json_response = json_format.MessageToDict(response)
-
-    vd = {}
-    vectors = json_response.get("vectors", {})
-    namespace = json_response.get("namespace", "")
-
-    for id, vec in vectors.items():
-        vd[id] = _Vector(
-            id=vec["id"],
-            values=vec.get("values", None),
-            sparse_values=parse_sparse_values(vec.get("sparseValues", None)),
-            metadata=vec.get("metadata", None),
-            _check_type=False,
-        )
-
-    pagination = None
-    if json_response.get("pagination") and json_response["pagination"].get("next"):
-        pagination = Pagination(next=json_response["pagination"]["next"])
+    """Parse a FetchByMetadataResponse protobuf message directly without MessageToDict conversion.
 
+    This optimized version directly accesses protobuf fields for better performance.
+    """
     # Extract response info from initial metadata
     from pinecone.utils.response_info import extract_response_info
 
     metadata = initial_metadata or {}
     response_info = extract_response_info(metadata)
 
+    # Directly access protobuf fields instead of converting entire message to dict
+    vd = {}
+    # namespace is a required string field, so it will always have a value (default empty string)
+    namespace = response.namespace
+
+    # Iterate over vectors map directly
+    for vec_id, vec in response.vectors.items():
+        # Convert vector.values (RepeatedScalarFieldContainer) to list
+        values = list(vec.values) if vec.values else None
+
+        # Handle sparse_values if present
+        parsed_sparse = None
+        if vec.HasField("sparse_values") and vec.sparse_values:
+            parsed_sparse = parse_sparse_values(
+                {
+                    "indices": list(vec.sparse_values.indices),
+                    "values": list(vec.sparse_values.values),
+                }
+            )
+
+        # Convert metadata Struct to dict only when needed
+        metadata_dict = None
+        if vec.HasField("metadata") and vec.metadata:
+            metadata_dict = json_format.MessageToDict(vec.metadata)
+
+        vd[vec_id] = _Vector(
+            id=vec.id,
+            values=values,
+            sparse_values=parsed_sparse,
+            metadata=metadata_dict,
+            _check_type=False,
+        )
+
+    # Parse pagination if present
+    pagination = None
+    if response.HasField("pagination") and response.pagination:
+        pagination = Pagination(next=response.pagination.next)
+
+    # Parse usage if present
     usage = None
-    if json_response.get("usage"):
-        usage = parse_usage(json_response.get("usage", {}))
+    if response.HasField("usage") and response.usage:
+        usage = parse_usage({"readUnits": response.usage.read_units})
+
     fetch_by_metadata_response = FetchByMetadataResponse(
         vectors=vd,
         namespace=namespace,
@@ -153,43 +191,50 @@ def parse_usage(usage: dict) -> Usage:
 
 
 def parse_upsert_response(
-    response: Message, _check_type: bool = False, initial_metadata: dict[str, str] | None = None
+    response: "ProtoUpsertResponse",
+    _check_type: bool = False,
+    initial_metadata: dict[str, str] | None = None,
 ) -> UpsertResponse:
-    from pinecone.utils.response_info import extract_response_info
+    """Parse an UpsertResponse protobuf message directly without MessageToDict conversion.
 
-    json_response = json_format.MessageToDict(response)
-    upserted_count = json_response.get("upsertedCount", 0)
+    This optimized version directly accesses protobuf fields for better performance.
+    """
+    from pinecone.utils.response_info import extract_response_info
 
     # Extract response info from initial metadata
     # For gRPC, LSN headers are in initial_metadata
     metadata = initial_metadata or {}
     response_info = extract_response_info(metadata)
 
+    # Directly access upserted_count field (required field in proto3, always has a value)
+    upserted_count = response.upserted_count
+
     return UpsertResponse(upserted_count=int(upserted_count), _response_info=response_info)
 
 
 def parse_update_response(
-    response: dict | Message,
+    response: dict | "ProtoUpdateResponse",
     _check_type: bool = False,
     initial_metadata: dict[str, str] | None = None,
 ) -> UpdateResponse:
+    """Parse an UpdateResponse protobuf message directly without MessageToDict conversion.
+
+    This optimized version directly accesses protobuf fields for better performance.
+    For dict responses (REST API), falls back to the original dict-based parsing.
+    """
     from pinecone.utils.response_info import extract_response_info
-    from google.protobuf import json_format
 
     # Extract response info from initial metadata
     metadata = initial_metadata or {}
     response_info = extract_response_info(metadata)
 
     # Extract matched_records from response
     matched_records = None
-    if isinstance(response, Message):
-        # GRPC response - convert to dict to extract matched_records
-        json_response = json_format.MessageToDict(response)
-        matched_records = json_response.get("matchedRecords") or json_response.get(
-            "matched_records"
-        )
+    if isinstance(response, Message) and not isinstance(response, dict):
+        # Optimized path: directly access protobuf field
+        matched_records = response.matched_records if response.HasField("matched_records") else None
     elif isinstance(response, dict):
-        # Dict response - extract directly
+        # Fallback for dict responses (REST API)
         matched_records = response.get("matchedRecords") or response.get("matched_records")
 
     return UpdateResponse(matched_records=matched_records, _response_info=response_info)
@@ -211,7 +256,7 @@ def parse_delete_response(
 
 
 def parse_query_response(
-    response: dict | Message,
+    response: dict | "ProtoQueryResponse",
     _check_type: bool = False,
     initial_metadata: dict[str, str] | None = None,
 ) -> QueryResponse:
@@ -226,7 +271,7 @@ def parse_query_response(
     metadata = initial_metadata or {}
     response_info = extract_response_info(metadata)
 
-    if isinstance(response, Message):
+    if isinstance(response, Message) and not isinstance(response, dict):
         # Optimized path: directly access protobuf fields
         matches = []
         # namespace is a required string field, so it will always have a value (default empty string)
@@ -320,26 +365,30 @@ def parse_stats_response(response: dict) -> "DescribeIndexStatsResponse":
 
 
 def parse_namespace_description(
-    response: Message, initial_metadata: dict[str, str] | None = None
+    response: "ProtoNamespaceDescription", initial_metadata: dict[str, str] | None = None
 ) -> NamespaceDescription:
+    """Parse a NamespaceDescription protobuf message directly without MessageToDict conversion.
+
+    This optimized version directly accesses protobuf fields for better performance.
+    """
     from pinecone.utils.response_info import extract_response_info
 
-    json_response = json_format.MessageToDict(response)
+    # Directly access protobuf fields
+    name = response.name
+    record_count = response.record_count
 
     # Extract indexed_fields if present
     indexed_fields = None
-    if "indexedFields" in json_response and json_response["indexedFields"]:
-        indexed_fields_data = json_response["indexedFields"]
-        if "fields" in indexed_fields_data:
+    if response.HasField("indexed_fields") and response.indexed_fields:
+        # Access indexed_fields.fields directly (RepeatedScalarFieldContainer)
+        fields_list = list(response.indexed_fields.fields) if response.indexed_fields.fields else []
+        if fields_list:
             indexed_fields = NamespaceDescriptionIndexedFields(
-                fields=indexed_fields_data.get("fields", []), _check_type=False
+                fields=fields_list, _check_type=False
             )
 
     namespace_desc = NamespaceDescription(
-        name=json_response.get("name", ""),
-        record_count=json_response.get("recordCount", 0),
-        indexed_fields=indexed_fields,
-        _check_type=False,
+        name=name, record_count=record_count, indexed_fields=indexed_fields, _check_type=False
     )
 
     # Attach _response_info as an attribute (NamespaceDescription is an OpenAPI model)
@@ -352,36 +401,44 @@ def parse_namespace_description(
     return cast(NamespaceDescription, namespace_desc)
 
 
-def parse_list_namespaces_response(response: Message) -> ListNamespacesResponse:
-    json_response = json_format.MessageToDict(response)
+def parse_list_namespaces_response(
+    response: "ProtoListNamespacesResponse",
+) -> ListNamespacesResponse:
+    """Parse a ListNamespacesResponse protobuf message directly without MessageToDict conversion.
 
+    This optimized version directly accesses protobuf fields for better performance.
+    """
+    # Directly iterate over namespaces
     namespaces = []
-    for ns in json_response.get("namespaces", []):
+    for ns in response.namespaces:
         # Extract indexed_fields if present
         indexed_fields = None
-        if "indexedFields" in ns and ns["indexedFields"]:
-            indexed_fields_data = ns["indexedFields"]
-            if "fields" in indexed_fields_data:
+        if ns.HasField("indexed_fields") and ns.indexed_fields:
+            # Access indexed_fields.fields directly (RepeatedScalarFieldContainer)
+            fields_list = list(ns.indexed_fields.fields) if ns.indexed_fields.fields else []
+            if fields_list:
                 indexed_fields = NamespaceDescriptionIndexedFields(
-                    fields=indexed_fields_data.get("fields", []), _check_type=False
+                    fields=fields_list, _check_type=False
                 )
 
         namespaces.append(
             NamespaceDescription(
-                name=ns.get("name", ""),
-                record_count=ns.get("recordCount", 0),
+                name=ns.name,
+                record_count=ns.record_count,
                 indexed_fields=indexed_fields,
                 _check_type=False,
             )
         )
 
+    # Parse pagination if present
     pagination = None
-    if "pagination" in json_response and json_response["pagination"]:
-        pagination = OpenApiPagination(
-            next=json_response["pagination"].get("next", ""), _check_type=False
-        )
+    if response.HasField("pagination") and response.pagination:
+        pagination = OpenApiPagination(next=response.pagination.next, _check_type=False)
+
+    # Parse total_count (int field in proto3, always has a value, default 0)
+    # If 0, treat as None to match original behavior
+    total_count = response.total_count if response.total_count else None
 
-    total_count = json_response.get("totalCount")
     from typing import cast
 
     result = ListNamespacesResponse(

tests/perf/test_grpc_parsing_perf.py

@@ -11,12 +11,22 @@
 from pinecone.core.grpc.protos.db_data_2025_10_pb2 import (
     FetchResponse,
     QueryResponse,
+    FetchByMetadataResponse,
+    UpsertResponse,
+    UpdateResponse,
     Vector,
     ScoredVector,
     SparseValues,
     Usage,
+    Pagination,
+)
+from pinecone.grpc.utils import (
+    parse_fetch_response,
+    parse_query_response,
+    parse_fetch_by_metadata_response,
+    parse_upsert_response,
+    parse_update_response,
 )
-from pinecone.grpc.utils import parse_fetch_response, parse_query_response
 
 
 def create_vector(id: str, dimension: int, include_sparse: bool = False) -> Vector:
@@ -104,6 +114,35 @@ def create_query_response(
     )
 
 
+def create_fetch_by_metadata_response(
+    num_vectors: int, dimension: int, include_sparse: bool = False
+) -> FetchByMetadataResponse:
+    """Create a FetchByMetadataResponse protobuf message with specified number of vectors."""
+    vectors = {}
+    for i in range(num_vectors):
+        vector = create_vector(f"vec_{i}", dimension, include_sparse)
+        vectors[f"vec_{i}"] = vector
+
+    pagination = Pagination(next="next_token") if num_vectors > 10 else None
+
+    return FetchByMetadataResponse(
+        vectors=vectors,
+        namespace="test_namespace",
+        usage=Usage(read_units=num_vectors),
+        pagination=pagination,
+    )
+
+
+def create_upsert_response(upserted_count: int) -> UpsertResponse:
+    """Create an UpsertResponse protobuf message."""
+    return UpsertResponse(upserted_count=upserted_count)
+
+
+def create_update_response(matched_records: int) -> UpdateResponse:
+    """Create an UpdateResponse protobuf message."""
+    return UpdateResponse(matched_records=matched_records)
+
+
 class TestFetchResponseParsingPerf:
     """Performance benchmarks for parse_fetch_response."""
 
@@ -160,3 +199,31 @@ def test_parse_query_response_sparse(self, benchmark, num_matches, dimension):
         """Benchmark parse_query_response with sparse vectors."""
         response = create_query_response(num_matches, dimension, include_sparse=True)
         benchmark(parse_query_response, response, False, None)
+
+
+class TestFetchByMetadataResponseParsingPerf:
+    """Performance benchmarks for parse_fetch_by_metadata_response."""
+
+    @pytest.mark.parametrize("num_vectors,dimension", [(10, 128), (100, 128), (1000, 128)])
+    def test_parse_fetch_by_metadata_response_dense(self, benchmark, num_vectors, dimension):
+        """Benchmark parse_fetch_by_metadata_response with dense vectors."""
+        response = create_fetch_by_metadata_response(num_vectors, dimension, include_sparse=False)
+        benchmark(parse_fetch_by_metadata_response, response, None)
+
+
+class TestUpsertResponseParsingPerf:
+    """Performance benchmarks for parse_upsert_response."""
+
+    def test_parse_upsert_response(self, benchmark):
+        """Benchmark parse_upsert_response."""
+        response = create_upsert_response(upserted_count=100)
+        benchmark(parse_upsert_response, response, False, None)
+
+
+class TestUpdateResponseParsingPerf:
+    """Performance benchmarks for parse_update_response."""
+
+    def test_parse_update_response(self, benchmark):
+        """Benchmark parse_update_response."""
+        response = create_update_response(matched_records=50)
+        benchmark(parse_update_response, response, False, None)