⚡️ Speed up function parse_query_response by 10%

[codeflash-ai]

📄 10% (0.10x) speedup for parse_query_response in pinecone/db_data/index.py

⏱️ Runtime : 21.9 microseconds → 19.9 microseconds (best of 434 runs)

📝 Explanation and details

The optimization replaces dict.pop("results", None) with a conditional check followed by del when the key exists. This provides a 10% speedup by eliminating unnecessary operations in two scenarios:

Key optimization changes:

1. Conditional existence check: if "results" in response._data_store: only performs deletion when the key exists
2. Direct deletion: del response._data_store["results"] is faster than pop() when you don't need the returned value

Why this is faster:

• The original pop("results", None) always performs an internal dictionary lookup, even when the key doesn't exist, then returns the default value
• The optimized version avoids the deletion operation entirely when "results" is absent, which is more efficient than pop()'s internal default value handling
• del is inherently faster than pop() because it doesn't need to return a value

Performance characteristics from tests:

• Best gains (20-40% faster): When "results" key is absent (test_basic_no_results_key: 41.3% faster, test_results_key_absent: 24.4% faster)
• Consistent gains (5-25% faster): When "results" key exists across various data types and sizes
• Minimal overhead: Even with large dictionaries (1000+ keys), maintains 5-12% improvement

This optimization is particularly effective for workloads where the "results" key is frequently absent, but provides consistent benefits regardless of the key's presence.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import pytest
from pinecone.db_data.index import parse_query_response


# function to test
# Simulate the QueryResponse class for testing purposes
class QueryResponse:
    def __init__(self, results=None, **kwargs):
        # Store results in _data_store under "results", and other kwargs as additional keys
        self._data_store = {}
        if results is not None:
            self._data_store["results"] = results
        for k, v in kwargs.items():
            self._data_store[k] = v

    # For equality testing, compare _data_store
    def __eq__(self, other):
        if not isinstance(other, QueryResponse):
            return False
        return self._data_store == other._data_store

    def __repr__(self):
        return f"QueryResponse({self._data_store})"
from pinecone.db_data.index import parse_query_response

# unit tests

# ----------- Basic Test Cases -----------

def test_basic_removes_results_key():
    # Test: results key is present and should be removed
    resp = QueryResponse(results=[1,2,3], foo="bar")
    codeflash_output = parse_query_response(resp); out = codeflash_output # 513ns -> 458ns (12.0% faster)

def test_basic_no_results_key():
    # Test: results key is absent; function should not fail
    resp = QueryResponse(foo="bar")
    codeflash_output = parse_query_response(resp); out = codeflash_output # 445ns -> 315ns (41.3% faster)

def test_basic_empty_results_list():
    # Test: results is an empty list; should still be removed
    resp = QueryResponse(results=[], foo="baz")
    codeflash_output = parse_query_response(resp); out = codeflash_output # 452ns -> 448ns (0.893% faster)

def test_basic_returns_same_object():
    # Test: Should return the same object (not a copy)
    resp = QueryResponse(results=[1,2,3])
    codeflash_output = parse_query_response(resp); out = codeflash_output # 435ns -> 430ns (1.16% faster)

# ----------- Edge Test Cases -----------

def test_edge_results_is_none():
    # Test: results is explicitly None; should be removed
    resp = QueryResponse(results=None, foo="edge")
    codeflash_output = parse_query_response(resp); out = codeflash_output # 434ns -> 345ns (25.8% faster)

def test_edge_results_is_falsey_value():
    # Test: results is a falsey value (0)
    resp = QueryResponse(results=0, foo="zero")
    codeflash_output = parse_query_response(resp); out = codeflash_output # 441ns -> 397ns (11.1% faster)

def test_edge_results_is_empty_dict():
    # Test: results is an empty dict
    resp = QueryResponse(results={}, foo="dict")
    codeflash_output = parse_query_response(resp); out = codeflash_output # 496ns -> 419ns (18.4% faster)

def test_edge_results_is_string():
    # Test: results is a string
    resp = QueryResponse(results="string", foo="bar")
    codeflash_output = parse_query_response(resp); out = codeflash_output # 467ns -> 409ns (14.2% faster)

def test_edge_results_is_nested_structure():
    # Test: results is a nested structure
    resp = QueryResponse(results=[{"a": [1,2,3]}, {"b": {"c": "d"}}], foo="nested")
    codeflash_output = parse_query_response(resp); out = codeflash_output # 661ns -> 614ns (7.65% faster)

def test_edge_multiple_other_keys():
    # Test: multiple other keys, results should be only one removed
    resp = QueryResponse(results=[1], a=1, b=2, c=3)
    codeflash_output = parse_query_response(resp); out = codeflash_output # 473ns -> 416ns (13.7% faster)

def test_edge_results_key_is_only_key():
    # Test: results is the only key
    resp = QueryResponse(results=[1,2,3])
    codeflash_output = parse_query_response(resp); out = codeflash_output # 491ns -> 460ns (6.74% faster)

def test_edge_results_is_unusual_type():
    # Test: results is an unusual type (tuple)
    resp = QueryResponse(results=(1,2,3), foo="tuple")
    codeflash_output = parse_query_response(resp); out = codeflash_output # 456ns -> 380ns (20.0% faster)

def test_edge_results_is_boolean():
    # Test: results is a boolean value
    resp = QueryResponse(results=True, foo="bool")
    codeflash_output = parse_query_response(resp); out = codeflash_output # 447ns -> 417ns (7.19% faster)

def test_edge_results_is_float():
    # Test: results is a float
    resp = QueryResponse(results=3.14, foo="float")
    codeflash_output = parse_query_response(resp); out = codeflash_output # 455ns -> 395ns (15.2% faster)

def test_edge_results_is_bytes():
    # Test: results is bytes
    resp = QueryResponse(results=b"bytes", foo="bytes")
    codeflash_output = parse_query_response(resp); out = codeflash_output # 426ns -> 388ns (9.79% faster)

# ----------- Large Scale Test Cases -----------

def test_large_many_other_keys():
    # Test: QueryResponse with many other keys
    keys = {f"key_{i}": i for i in range(1000)}
    resp = QueryResponse(results=list(range(100)), **keys)
    codeflash_output = parse_query_response(resp); out = codeflash_output # 717ns -> 680ns (5.44% faster)
    for k in keys:
        pass

def test_large_results_list():
    # Test: results is a large list
    large_list = list(range(1000))
    resp = QueryResponse(results=large_list, foo="large")
    codeflash_output = parse_query_response(resp); out = codeflash_output # 517ns -> 452ns (14.4% faster)

def test_large_results_large_nested_structure():
    # Test: results is a large nested structure
    nested = [{"a": [i, i+1, i+2]} for i in range(1000)]
    resp = QueryResponse(results=nested, foo="nested")
    codeflash_output = parse_query_response(resp); out = codeflash_output # 549ns -> 486ns (13.0% faster)

def test_large_empty_results_and_many_keys():
    # Test: results is empty, many other keys
    keys = {f"key_{i}": i for i in range(900)}
    resp = QueryResponse(results=[], **keys)
    codeflash_output = parse_query_response(resp); out = codeflash_output # 606ns -> 576ns (5.21% faster)
    for k in keys:
        pass

def test_large_multiple_calls_idempotency():
    # Test: Calling parse_query_response multiple times should not fail
    resp = QueryResponse(results=[1,2,3], foo="bar")
    codeflash_output = parse_query_response(resp); out1 = codeflash_output # 508ns -> 522ns (2.68% slower)
    codeflash_output = parse_query_response(out1); out2 = codeflash_output # 257ns -> 255ns (0.784% faster)

# ----------- Mutation Testing: Defensive/Negative -----------

def test_mutation_wrong_key_removed():
    # If function removes a wrong key, test should fail
    resp = QueryResponse(results=[1,2,3], foo="bar")
    # Simulate mutation: remove "foo" instead of "results"
    resp._data_store.pop("foo", None)

def test_mutation_returns_new_object():
    # If function returns a new object instead of modifying in-place, test should fail
    resp = QueryResponse(results=[1,2,3], foo="bar")
    codeflash_output = parse_query_response(resp); out = codeflash_output # 538ns -> 511ns (5.28% faster)

def test_mutation_does_not_remove_results():
    # If function does not remove "results", test should fail
    resp = QueryResponse(results=[1,2,3], foo="bar")
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
import copy
from types import SimpleNamespace

# imports
import pytest  # used for our unit tests
from pinecone.db_data.index import parse_query_response


# function to test
# Simulate QueryResponse for testing purposes
class QueryResponse:
    def __init__(self, results=None, **kwargs):
        # Store all passed kwargs in _data_store for mutation
        self._data_store = dict(kwargs)
        if results is not None:
            self._data_store["results"] = results
    def __eq__(self, other):
        # Equality based on _data_store contents
        return isinstance(other, QueryResponse) and self._data_store == other._data_store
    def __repr__(self):
        return f"QueryResponse({self._data_store})"
from pinecone.db_data.index import parse_query_response

# unit tests

# 1. Basic Test Cases

def test_removes_results_key_when_present():
    # Should remove 'results' if present
    resp = QueryResponse(results=[1,2,3], foo="bar")
    codeflash_output = parse_query_response(resp); out = codeflash_output # 610ns -> 545ns (11.9% faster)

def test_returns_same_object():
    # Should return the same object, not a copy
    resp = QueryResponse(results=[42])
    codeflash_output = parse_query_response(resp); out = codeflash_output # 561ns -> 493ns (13.8% faster)

def test_does_not_modify_other_keys():
    # Should not affect other keys
    resp = QueryResponse(results=[1], a=1, b=2)
    codeflash_output = parse_query_response(resp); out = codeflash_output # 500ns -> 477ns (4.82% faster)

def test_results_key_absent():
    # Should not fail if 'results' is missing
    resp = QueryResponse(foo="bar")
    before = copy.deepcopy(resp._data_store)
    codeflash_output = parse_query_response(resp); out = codeflash_output # 510ns -> 410ns (24.4% faster)

def test_results_key_none():
    # Should remove 'results' even if value is None
    resp = QueryResponse(results=None, foo="bar")
    codeflash_output = parse_query_response(resp); out = codeflash_output # 415ns -> 375ns (10.7% faster)

def test_results_key_empty_list():
    # Should remove 'results' if it's an empty list
    resp = QueryResponse(results=[], foo="bar")
    codeflash_output = parse_query_response(resp); out = codeflash_output # 477ns -> 445ns (7.19% faster)

# 2. Edge Test Cases

def test_results_key_is_unusual_type():
    # Should remove 'results' regardless of value type
    resp = QueryResponse(results={"unexpected": "dict"}, foo="bar")
    codeflash_output = parse_query_response(resp); out = codeflash_output # 482ns -> 452ns (6.64% faster)

def test_other_keys_with_similar_names():
    # Should not remove keys with similar names
    resp = QueryResponse(results=[1], results_extra="keep_me", foo="bar")
    codeflash_output = parse_query_response(resp); out = codeflash_output # 500ns -> 472ns (5.93% faster)

def test_multiple_removals():
    # Should not error if called multiple times
    resp = QueryResponse(results=[1], foo="bar")
    parse_query_response(resp) # 493ns -> 400ns (23.2% faster)
    codeflash_output = parse_query_response(resp); out = codeflash_output # 260ns -> 258ns (0.775% faster)

def test_empty_data_store():
    # Should not fail if _data_store is empty
    resp = QueryResponse()
    codeflash_output = parse_query_response(resp); out = codeflash_output # 433ns -> 377ns (14.9% faster)

def test_results_key_is_falsey():
    # Should remove 'results' if it's False, 0, or ""
    for val in [False, 0, ""]:
        resp = QueryResponse(results=val, foo="bar")
        codeflash_output = parse_query_response(resp); out = codeflash_output # 803ns -> 788ns (1.90% faster)

def test_results_key_is_tuple():
    # Should remove 'results' if it's a tuple
    resp = QueryResponse(results=(1,2,3), foo="bar")
    codeflash_output = parse_query_response(resp); out = codeflash_output # 435ns -> 362ns (20.2% faster)

def test_results_key_is_set():
    # Should remove 'results' if it's a set
    resp = QueryResponse(results={1,2,3}, foo="bar")
    codeflash_output = parse_query_response(resp); out = codeflash_output # 510ns -> 436ns (17.0% faster)

def test_results_key_is_bytes():
    # Should remove 'results' if it's bytes
    resp = QueryResponse(results=b"123", foo="bar")
    codeflash_output = parse_query_response(resp); out = codeflash_output # 443ns -> 402ns (10.2% faster)

def test_results_key_is_nested():
    # Should remove 'results' if it's a nested structure
    resp = QueryResponse(results=[{"a":1}, {"b":2}], foo="bar")
    codeflash_output = parse_query_response(resp); out = codeflash_output # 560ns -> 463ns (21.0% faster)

def test_results_key_is_object():
    # Should remove 'results' if it's an object
    obj = SimpleNamespace(a=1)
    resp = QueryResponse(results=obj, foo="bar")
    codeflash_output = parse_query_response(resp); out = codeflash_output # 414ns -> 386ns (7.25% faster)

# 3. Large Scale Test Cases

def test_large_number_of_other_keys():
    # Should only remove 'results' even with many other keys
    keys = {f"key_{i}": i for i in range(999)}
    resp = QueryResponse(results=[1,2,3], **keys)
    codeflash_output = parse_query_response(resp); out = codeflash_output # 645ns -> 575ns (12.2% faster)
    for k, v in keys.items():
        pass

def test_large_results_list():
    # Should remove 'results' even if it's a large list
    large_list = list(range(1000))
    resp = QueryResponse(results=large_list, foo="bar")
    codeflash_output = parse_query_response(resp); out = codeflash_output # 537ns -> 507ns (5.92% faster)

def test_large_nested_results():
    # Should remove 'results' with large nested structure
    nested = [{"a": [i for i in range(10)]} for _ in range(100)]
    resp = QueryResponse(results=nested, foo="bar")
    codeflash_output = parse_query_response(resp); out = codeflash_output # 476ns -> 455ns (4.62% faster)

def test_performance_large_keys_and_results():
    # Should perform efficiently with large keys and results
    keys = {f"key_{i}": i for i in range(500)}
    large_list = list(range(500))
    resp = QueryResponse(results=large_list, **keys)
    codeflash_output = parse_query_response(resp); out = codeflash_output # 508ns -> 540ns (5.93% slower)
    for k, v in keys.items():
        pass

def test_no_side_effects_on_other_instances():
    # Should not affect other QueryResponse instances
    resp1 = QueryResponse(results=[1,2,3], foo="bar")
    resp2 = QueryResponse(results=[4,5,6], foo="baz")
    parse_query_response(resp1) # 526ns -> 472ns (11.4% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-parse_query_response-mh6c0sl8 and push.