⚡️ Speed up method GrpcRunner._prepare_metadata by 58%

[codeflash-ai]

📄 58% (0.58x) speedup for GrpcRunner._prepare_metadata in pinecone/grpc/grpc_runner.py

⏱️ Runtime : 759 microseconds → 480 microseconds (best of 217 runs)

📝 Explanation and details

The optimization achieves a 58% speedup by eliminating expensive dictionary unpacking operations and reducing allocations during metadata preparation.

Key optimizations:

1. Pre-computed fixed metadata: The constructor now stores fixed metadata as both a list of tuples (_fixed_metadata_items) and a dictionary (fixed_metadata), avoiding repeated dictionary creation and merging operations.

2. Eliminated dictionary unpacking: The original code used expensive dictionary unpacking (**self.fixed_metadata, **self._request_metadata(), **user_provided_metadata) which creates multiple intermediate dictionaries. The optimized version uses dict.copy() and dict.update() operations instead.

3. Conditional user metadata handling: Added a check if user_provided_metadata: to avoid unnecessary update operations when no user metadata is provided.

Why this is faster:

• Dictionary unpacking with ** operator is expensive in Python as it creates new dictionary objects for each unpacking operation
• dict.copy() followed by dict.update() is significantly faster than unpacking multiple dictionaries
• Pre-computing fixed metadata eliminates redundant work on every call

Performance benefits by test case:

• Basic cases: 11-19% faster for small metadata sets
• Large metadata: 61-119% faster - the optimization scales particularly well with larger datasets (500-1000 metadata entries)
• Edge cases: Consistent 11-19% improvements across various edge conditions

The optimization is especially effective for applications with frequent metadata preparation calls and large amounts of additional/user metadata, which is common in gRPC service implementations.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 60 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Generated Regression Tests and Runtime

import uuid

# imports
import pytest  # used for our unit tests
from pinecone.grpc.grpc_runner import GrpcRunner


# Mocks for config/constants used in GrpcRunner
class Config:
    def __init__(self, api_key):
        self.api_key = api_key

class GRPCClientConfig:
    def __init__(self, additional_metadata=None):
        self.additional_metadata = additional_metadata or {}

CLIENT_VERSION = "1.2.3"
API_VERSION = "2023-12-01"
REQUEST_ID = "x-request-id"

# -------------------------------
# Unit Tests for _prepare_metadata
# -------------------------------

# BASIC TEST CASES

def test_basic_metadata_merging():
    """Basic: User metadata is merged with fixed and request metadata."""
    config = Config(api_key="test-key")
    grpc_config = GRPCClientConfig()
    runner = GrpcRunner("my_index", config, grpc_config)
    user_metadata = {"foo": "bar", "baz": "qux"}
    result = dict(runner._prepare_metadata(user_metadata))
    uuid.UUID(result[REQUEST_ID])  # Should not raise

def test_basic_no_user_metadata():
    """Basic: No user metadata, only fixed and request metadata."""
    config = Config(api_key="key123")
    grpc_config = GRPCClientConfig()
    runner = GrpcRunner("idx", config, grpc_config)
    result = dict(runner._prepare_metadata({}))
    uuid.UUID(result[REQUEST_ID])  # Should not raise

def test_basic_additional_metadata():
    """Basic: Additional metadata from GRPCClientConfig is merged."""
    config = Config(api_key="k")
    grpc_config = GRPCClientConfig(additional_metadata={"extra": "meta"})
    runner = GrpcRunner("i", config, grpc_config)
    result = dict(runner._prepare_metadata({})) # 14.5μs -> 13.0μs (11.6% faster)

def test_user_metadata_overrides_fixed_and_additional():
    """Basic: User metadata overrides fixed/additional metadata if keys overlap."""
    config = Config(api_key="fixed-key")
    grpc_config = GRPCClientConfig(additional_metadata={"api-key": "extra-key", "foo": "bar"})
    runner = GrpcRunner("index", config, grpc_config)
    user_metadata = {"api-key": "user-key", "foo": "user-bar"}
    result = dict(runner._prepare_metadata(user_metadata)
    ) # 11.4μs -> 10.2μs (11.6% faster)

def test_request_id_is_unique_each_call():
    """Basic: REQUEST_ID is different for each call."""
    config = Config(api_key="k")
    grpc_config = GRPCClientConfig()
    runner = GrpcRunner("i", config, grpc_config)
    id1 = dict(runner._prepare_metadata({}))[REQUEST_ID]
    id2 = dict(runner._prepare_metadata({}))[REQUEST_ID]
    uuid.UUID(id1)
    uuid.UUID(id2)

# EDGE TEST CASES

def test_empty_strings_and_special_characters():
    """Edge: Metadata keys/values can be empty or have special characters."""
    config = Config(api_key="")
    grpc_config = GRPCClientConfig(additional_metadata={"": "", "weird@key": "weird:value"})
    runner = GrpcRunner("", config, grpc_config)
    user_metadata = {"": "user_empty", "weird@key": "user_weird"}
    result = dict(runner._prepare_metadata(user_metadata)) # 14.9μs -> 13.1μs (13.2% faster)

def test_user_metadata_with_none_values():
    """Edge: None values in user metadata are preserved as is."""
    config = Config(api_key="none-key")
    grpc_config = GRPCClientConfig()
    runner = GrpcRunner("none-index", config, grpc_config)
    user_metadata = {"foo": None}
    result = dict(runner._prepare_metadata(user_metadata)) # 11.5μs -> 10.0μs (14.2% faster)

def test_overlapping_keys_all_sources():
    """Edge: Key overlaps between fixed, additional, and user metadata."""
    config = Config(api_key="fixed")
    grpc_config = GRPCClientConfig(additional_metadata={"service-name": "additional", "client-version": "additional-v"})
    runner = GrpcRunner("fixed-index", config, grpc_config)
    user_metadata = {"service-name": "user", "client-version": "user-v"}
    result = dict(runner._prepare_metadata(user_metadata)) # 10.8μs -> 9.36μs (15.3% faster)

def test_metadata_with_non_string_keys_and_values():
    """Edge: Non-string keys/values are converted to string in output tuple."""
    config = Config(api_key=123)
    grpc_config = GRPCClientConfig(additional_metadata={456: 789})
    runner = GrpcRunner(101112, config, grpc_config)
    user_metadata = {131415: 161718}
    result = dict(runner._prepare_metadata(user_metadata)) # 10.8μs -> 9.16μs (18.4% faster)

def test_metadata_with_reserved_request_id_key():
    """Edge: User metadata can override REQUEST_ID."""
    config = Config(api_key="k")
    grpc_config = GRPCClientConfig()
    runner = GrpcRunner("i", config, grpc_config)
    user_metadata = {REQUEST_ID: "user-id"}
    result = dict(runner._prepare_metadata(user_metadata)) # 10.3μs -> 9.19μs (12.5% faster)

# LARGE SCALE TEST CASES

def test_large_number_of_user_metadata_keys():
    """Large: Many user metadata keys."""
    config = Config(api_key="big-key")
    grpc_config = GRPCClientConfig()
    runner = GrpcRunner("big-index", config, grpc_config)
    user_metadata = {f"key{i}": f"value{i}" for i in range(500)}
    result = dict(runner._prepare_metadata(user_metadata)) # 38.6μs -> 22.4μs (72.4% faster)
    for i in range(500):
        pass

def test_large_additional_metadata():
    """Large: Many additional metadata keys."""
    config = Config(api_key="add-key")
    additional = {f"add{i}": f"meta{i}" for i in range(500)}
    grpc_config = GRPCClientConfig(additional_metadata=additional)
    runner = GrpcRunner("add-index", config, grpc_config)
    result = dict(runner._prepare_metadata({})) # 31.8μs -> 16.3μs (94.8% faster)
    for i in range(500):
        pass

def test_large_combined_metadata():
    """Large: Many keys in both user and additional metadata, with overlaps."""
    config = Config(api_key="combo-key")
    additional = {f"key{i}": f"add{i}" for i in range(250)}
    user_metadata = {f"key{i}": f"user{i}" for i in range(250, 500)}
    # Overlap: keys 0-249 in additional, 250-499 in user
    grpc_config = GRPCClientConfig(additional_metadata=additional)
    runner = GrpcRunner("combo-index", config, grpc_config)
    result = dict(runner._prepare_metadata(user_metadata)) # 37.5μs -> 21.9μs (71.5% faster)
    for i in range(250):
        pass
    for i in range(250, 500):
        pass

def test_large_metadata_user_overrides_fixed_and_additional():
    """Large: User metadata overrides fixed/additional for all keys."""
    config = Config(api_key="fixed-key")
    additional = {f"key{i}": f"add{i}" for i in range(500)}
    grpc_config = GRPCClientConfig(additional_metadata=additional)
    runner = GrpcRunner("index", config, grpc_config)
    user_metadata = {f"key{i}": f"user{i}" for i in range(500)}
    result = dict(runner._prepare_metadata(user_metadata)) # 40.8μs -> 25.2μs (61.7% faster)
    for i in range(500):
        pass

def test_large_metadata_performance():
    """Large: Performance test for large metadata (just ensure it runs quickly and correctly)."""
    config = Config(api_key="perf-key")
    additional = {f"add{i}": f"meta{i}" for i in range(999)}
    grpc_config = GRPCClientConfig(additional_metadata=additional)
    runner = GrpcRunner("perf-index", config, grpc_config)
    user_metadata = {f"user{i}": f"val{i}" for i in range(999)}
    # Should not raise or take excessive time
    result = dict(runner._prepare_metadata(user_metadata)) # 115μs -> 62.8μs (83.8% faster)
    for i in range(999):
        pass
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
import uuid
from typing import Dict, Tuple

# imports
import pytest  # used for our unit tests
from pinecone.grpc.grpc_runner import GrpcRunner

# --- Minimal implementation of dependencies for self-contained tests ---

CLIENT_VERSION = "1.2.3"
API_VERSION = "2023-06-01"
REQUEST_ID = "request-id"

class Config:
    def __init__(self, api_key: str):
        self.api_key = api_key

class GRPCClientConfig:
    def __init__(self, additional_metadata: Dict[str, str] = None):
        self.additional_metadata = additional_metadata or {}

# --- Unit tests ---

# Helper functions for test assertions

def _tuple_to_dict(metadata_tuple):
    """Convert metadata tuple to dict for easier assertions."""
    return dict(metadata_tuple)

def _is_valid_uuid(val):
    try:
        uuid.UUID(val)
        return True
    except Exception:
        return False

# --- Basic Test Cases ---

def test_basic_metadata_inclusion():
    """Test that fixed metadata and user metadata are included."""
    config = Config(api_key="test-key")
    grpc_config = GRPCClientConfig()
    runner = GrpcRunner(index_name="my-index", config=config, grpc_config=grpc_config)

    user_md = {"foo": "bar"}
    codeflash_output = runner._prepare_metadata(user_md); metadata = codeflash_output # 12.4μs -> 11.0μs (12.7% faster)
    md_dict = _tuple_to_dict(metadata)

def test_user_metadata_overrides_fixed():
    """Test that user metadata can override fixed metadata."""
    config = Config(api_key="fixed-key")
    grpc_config = GRPCClientConfig()
    runner = GrpcRunner(index_name="index", config=config, grpc_config=grpc_config)

    user_md = {
        "api-key": "user-key",
        "service-name": "user-index",
        "client-version": "user-version",
        "x-pinecone-api-version": "user-api-version"
    }
    codeflash_output = runner._prepare_metadata(user_md); metadata = codeflash_output # 10.7μs -> 9.42μs (13.5% faster)
    md_dict = _tuple_to_dict(metadata)

def test_additional_metadata_merging():
    """Test that additional_metadata in GRPCClientConfig is merged."""
    config = Config(api_key="test-key")
    additional = {"extra-header": "extra-value", "api-key": "overridden"}
    grpc_config = GRPCClientConfig(additional_metadata=additional)
    runner = GrpcRunner(index_name="idx", config=config, grpc_config=grpc_config)

    user_md = {"foo": "bar"}
    codeflash_output = runner._prepare_metadata(user_md); metadata = codeflash_output # 10.5μs -> 9.21μs (14.1% faster)
    md_dict = _tuple_to_dict(metadata)

def test_empty_user_metadata():
    """Test that function works with empty user metadata."""
    config = Config(api_key="test-key")
    grpc_config = GRPCClientConfig()
    runner = GrpcRunner(index_name="idx", config=config, grpc_config=grpc_config)

    codeflash_output = runner._prepare_metadata({}); metadata = codeflash_output # 10.1μs -> 8.59μs (17.6% faster)
    md_dict = _tuple_to_dict(metadata)

# --- Edge Test Cases ---

def test_empty_additional_metadata():
    """Test with empty additional_metadata in GRPCClientConfig."""
    config = Config(api_key="test-key")
    grpc_config = GRPCClientConfig(additional_metadata={})
    runner = GrpcRunner(index_name="idx", config=config, grpc_config=grpc_config)

    codeflash_output = runner._prepare_metadata({"foo": "bar"}); metadata = codeflash_output # 10.3μs -> 8.99μs (14.7% faster)
    md_dict = _tuple_to_dict(metadata)

def test_user_metadata_empty_string_key_and_value():
    """Test user metadata with empty string key and value."""
    config = Config(api_key="test-key")
    grpc_config = GRPCClientConfig()
    runner = GrpcRunner(index_name="idx", config=config, grpc_config=grpc_config)

    user_md = {"": ""}
    codeflash_output = runner._prepare_metadata(user_md); metadata = codeflash_output # 10.4μs -> 8.78μs (17.9% faster)
    md_dict = _tuple_to_dict(metadata)

def test_user_metadata_overrides_additional_and_fixed():
    """Test user metadata overrides additional_metadata and fixed_metadata."""
    config = Config(api_key="fixed-key")
    additional = {"api-key": "additional-key"}
    grpc_config = GRPCClientConfig(additional_metadata=additional)
    runner = GrpcRunner(index_name="idx", config=config, grpc_config=grpc_config)

    user_md = {"api-key": "user-key"}
    codeflash_output = runner._prepare_metadata(user_md); metadata = codeflash_output # 10.2μs -> 8.81μs (16.1% faster)
    md_dict = _tuple_to_dict(metadata)

def test_metadata_key_collision_order():
    """Test that user metadata overrides additional_metadata, which overrides fixed_metadata."""
    config = Config(api_key="fixed-key")
    additional = {"api-key": "additional-key"}
    grpc_config = GRPCClientConfig(additional_metadata=additional)
    runner = GrpcRunner(index_name="idx", config=config, grpc_config=grpc_config)

    # User metadata should override additional and fixed
    user_md = {"api-key": "user-key"}
    codeflash_output = runner._prepare_metadata(user_md); metadata = codeflash_output # 10.0μs -> 8.45μs (18.8% faster)
    md_dict = _tuple_to_dict(metadata)

    # If no user metadata, additional_metadata should override fixed
    codeflash_output = runner._prepare_metadata({}); metadata2 = codeflash_output # 4.91μs -> 4.13μs (19.0% faster)
    md_dict2 = _tuple_to_dict(metadata2)

def test_request_id_is_unique_per_call():
    """Test that request-id is different on each call."""
    config = Config(api_key="test-key")
    grpc_config = GRPCClientConfig()
    runner = GrpcRunner(index_name="idx", config=config, grpc_config=grpc_config)

    ids = set()
    for _ in range(5):
        codeflash_output = runner._prepare_metadata({}); metadata = codeflash_output
        md_dict = _tuple_to_dict(metadata)
        ids.add(md_dict[REQUEST_ID])
    for rid in ids:
        pass

def test_non_string_metadata_values():
    """Test that non-string values in user metadata are handled (should be converted to str)."""
    config = Config(api_key="test-key")
    grpc_config = GRPCClientConfig()
    runner = GrpcRunner(index_name="idx", config=config, grpc_config=grpc_config)

    user_md = {"int": 123, "float": 4.5, "bool": True, "none": None}
    codeflash_output = runner._prepare_metadata(user_md); metadata = codeflash_output # 14.7μs -> 13.2μs (11.1% faster)
    md_dict = _tuple_to_dict(metadata)

# Patch _prepare_metadata to coerce values to str for this test only

def test_large_user_metadata():
    """Test with a large number of user metadata entries."""
    config = Config(api_key="test-key")
    grpc_config = GRPCClientConfig()
    runner = GrpcRunner(index_name="idx", config=config, grpc_config=grpc_config)

    user_md = {f"key{i}": f"value{i}" for i in range(1000)}
    codeflash_output = runner._prepare_metadata(user_md); metadata = codeflash_output # 66.1μs -> 37.8μs (75.2% faster)
    md_dict = _tuple_to_dict(metadata)
    # All user keys should be present
    for i in range(1000):
        pass

def test_large_additional_metadata():
    """Test with a large number of additional_metadata entries."""
    config = Config(api_key="test-key")
    additional = {f"extra{i}": f"val{i}" for i in range(1000)}
    grpc_config = GRPCClientConfig(additional_metadata=additional)
    runner = GrpcRunner(index_name="idx", config=config, grpc_config=grpc_config)

    user_md = {}
    codeflash_output = runner._prepare_metadata(user_md); metadata = codeflash_output # 53.0μs -> 24.2μs (119% faster)
    md_dict = _tuple_to_dict(metadata)
    # All additional keys should be present
    for i in range(1000):
        pass

def test_large_metadata_key_collisions():
    """Test with large user metadata overriding large additional_metadata."""
    config = Config(api_key="test-key")
    additional = {f"key{i}": f"additional{i}" for i in range(1000)}
    grpc_config = GRPCClientConfig(additional_metadata=additional)
    runner = GrpcRunner(index_name="idx", config=config, grpc_config=grpc_config)

    user_md = {f"key{i}": f"user{i}" for i in range(1000)}
    codeflash_output = runner._prepare_metadata(user_md); metadata = codeflash_output # 70.9μs -> 42.7μs (65.9% faster)
    md_dict = _tuple_to_dict(metadata)
    # User metadata should override additional_metadata
    for i in range(1000):
        pass

def test_large_metadata_performance():
    """Test that function completes quickly with large metadata."""
    import time
    config = Config(api_key="test-key")
    additional = {f"extra{i}": f"val{i}" for i in range(1000)}
    grpc_config = GRPCClientConfig(additional_metadata=additional)
    runner = GrpcRunner(index_name="idx", config=config, grpc_config=grpc_config)
    user_md = {f"key{i}": f"value{i}" for i in range(1000)}

    start = time.time()
    codeflash_output = runner._prepare_metadata(user_md); metadata = codeflash_output # 116μs -> 61.6μs (88.8% faster)
    duration = time.time() - start
    # Spot check a few keys
    md_dict = _tuple_to_dict(metadata)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from pinecone.config.config import Config
from pinecone.grpc.config import GRPCClientConfig
from pinecone.grpc.grpc_runner import GrpcRunner

def test_GrpcRunner__prepare_metadata():
    GrpcRunner._prepare_metadata(GrpcRunner('', Config(api_key='', host='', proxy_url='', proxy_headers={}, ssl_ca_certs='', ssl_verify=False, additional_headers={}, source_tag=''), GRPCClientConfig(secure=True, timeout=0, conn_timeout=0, reuse_channel=False, retry_config=None, grpc_channel_options={}, additional_metadata=None)), {})

To edit these changes git checkout codeflash/optimize-GrpcRunner._prepare_metadata-mh9rhzdv and push.