⚡️ Speed up method Search.select_all by 107%

[codeflash-ai]

📄 107% (1.07x) speedup for Search.select_all in chromadb/execution/expression/plan.py

⏱️ Runtime : 1.98 milliseconds → 956 microseconds (best of 65 runs)

📝 Explanation and details

The optimization pre-computes the expensive Select object creation at module-level as _PREDEFINED_SELECT_ALL, eliminating the need to recreate it on every select_all() call.

Key changes:

• Module-level caching: _PREDEFINED_SELECT_ALL = Select(keys={Key.DOCUMENT, Key.EMBEDDING, Key.METADATA, Key.SCORE}) is created once when the module loads instead of every time select_all() is called
• Direct reference: The select_all() method now directly uses the cached object instead of constructing a new one

Why this creates a 107% speedup:
The line profiler shows the original code spent 48.3% of its time (3.95ms out of 8.18ms) creating the Select object with the set of four Key objects. Set construction and Key object creation are expensive operations in Python. By moving this to module load time, each select_all() call now only performs the much cheaper Search constructor call.

Test case performance patterns:

• Best for frequent calls: Tests like test_select_all_performance_large_loop show dramatic improvements (108% speedup) when select_all() is called repeatedly
• Consistent across scenarios: All test cases show 70-150% speedups regardless of the initial Search configuration, indicating the optimization benefits any usage pattern
• Maintains correctness: The cached object is immutable, so sharing it across calls is safe and doesn't affect functionality

The optimization is particularly effective because select_all() always returns the same set of keys, making it a perfect candidate for pre-computation.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import pytest
from chromadb.execution.expression.plan import Search

# --- Minimal stubs for Key, Select, and Search to enable testing ---

class Key:
    def __init__(self, name):
        self.name = name

    def __eq__(self, other):
        if isinstance(other, Key):
            return self.name == other.name
        return False

    def __hash__(self):
        return hash(self.name)

    def __repr__(self):
        return f"Key({self.name!r})"

# Predefined keys
Key.ID = Key("#id")
Key.DOCUMENT = Key("#document")
Key.EMBEDDING = Key("#embedding")
Key.METADATA = Key("#metadata")
Key.SCORE = Key("#score")

class Select:
    def __init__(self, keys=None):
        if keys is None:
            self.keys = set()
        else:
            self.keys = set(keys)

    @staticmethod
    def from_dict(data):
        keys = data.get("keys", [])
        key_objs = set()
        for k in keys:
            if k == "#id":
                key_objs.add(Key.ID)
            elif k == "#document":
                key_objs.add(Key.DOCUMENT)
            elif k == "#embedding":
                key_objs.add(Key.EMBEDDING)
            elif k == "#metadata":
                key_objs.add(Key.METADATA)
            elif k == "#score":
                key_objs.add(Key.SCORE)
            else:
                key_objs.add(Key(k))
        return Select(keys=key_objs)

    def __eq__(self, other):
        if not isinstance(other, Select):
            return False
        return self.keys == other.keys

    def __repr__(self):
        return f"Select(keys={self.keys!r})"

class Limit:
    def __init__(self, limit=None, offset=0):
        self.limit = limit
        self.offset = offset

    @staticmethod
    def from_dict(data):
        return Limit(limit=data.get("limit"), offset=data.get("offset", 0))

    def __eq__(self, other):
        return isinstance(other, Limit) and self.limit == other.limit and self.offset == other.offset

    def __repr__(self):
        return f"Limit(limit={self.limit!r}, offset={self.offset!r})"
from chromadb.execution.expression.plan import Search

# --- Unit Tests ---

# Helper: the set of all keys that select_all should select
ALL_KEYS_SET = {Key.DOCUMENT, Key.EMBEDDING, Key.METADATA, Key.SCORE}

# -------- BASIC TEST CASES --------

def test_select_all_returns_search_instance():
    # Should return a new Search object
    s = Search()
    codeflash_output = s.select_all(); s2 = codeflash_output # 3.24μs -> 1.47μs (121% faster)

def test_select_all_selects_all_predefined_keys():
    # Should select all four keys
    s = Search()
    codeflash_output = s.select_all(); s2 = codeflash_output # 2.81μs -> 1.42μs (97.7% faster)




def test_select_all_with_none_fields():
    # Should work even if all fields are None
    s = Search(where=None, rank=None, limit=None, select=None)
    codeflash_output = s.select_all(); s2 = codeflash_output # 3.52μs -> 1.40μs (151% faster)

def test_select_all_with_different_limit_types():
    # Should preserve int, dict, and Limit types for limit
    s1 = Search(limit=7)
    codeflash_output = s1.select_all(); s2 = codeflash_output

    s3 = Search(limit={"limit": 5, "offset": 2})
    codeflash_output = s3.select_all(); s4 = codeflash_output

    lim = Limit(limit=3, offset=1)
    s5 = Search(limit=lim)
    codeflash_output = s5.select_all(); s6 = codeflash_output

def test_select_all_with_select_as_list_or_set():
    # Should override select even if it was passed as a list or set
    s1 = Search(select=["#document", "#score"])
    codeflash_output = s1.select_all(); s2 = codeflash_output # 2.85μs -> 1.64μs (73.6% faster)

    s3 = Search(select={"#embedding", "#metadata"})
    codeflash_output = s3.select_all(); s4 = codeflash_output # 1.35μs -> 677ns (98.8% faster)








def test_select_all_does_not_include_id_key():
    # select_all should NOT include Key.ID
    codeflash_output = Search().select_all(); s = codeflash_output # 3.77μs -> 1.73μs (118% faster)


#------------------------------------------------
import pytest
from chromadb.execution.expression.plan import Search


# Minimal stubs for Key, Select, Search to enable testing
class Key:
    def __init__(self, name):
        self.name = name

    def __eq__(self, other):
        if isinstance(other, Key):
            return self.name == other.name
        return False

    def __hash__(self):
        return hash(self.name)

    def __repr__(self):
        return f"Key({self.name!r})"

# Predefined keys
Key.ID = Key("#id")
Key.DOCUMENT = Key("#document")
Key.EMBEDDING = Key("#embedding")
Key.METADATA = Key("#metadata")
Key.SCORE = Key("#score")

class Select:
    def __init__(self, keys=None):
        if keys is None:
            self.keys = set()
        else:
            self.keys = set(keys)

    @staticmethod
    def from_dict(data):
        keys = data.get("keys", [])
        key_objs = []
        for k in keys:
            if k == "#id":
                key_objs.append(Key.ID)
            elif k == "#document":
                key_objs.append(Key.DOCUMENT)
            elif k == "#embedding":
                key_objs.append(Key.EMBEDDING)
            elif k == "#metadata":
                key_objs.append(Key.METADATA)
            elif k == "#score":
                key_objs.append(Key.SCORE)
            else:
                key_objs.append(Key(k))
        return Select(keys=key_objs)

    def __eq__(self, other):
        return isinstance(other, Select) and self.keys == other.keys

    def __repr__(self):
        return f"Select(keys={self.keys!r})"
from chromadb.execution.expression.plan import Search

# Helper for test: set of all keys that select_all should select
ALL_KEYS = {Key.DOCUMENT, Key.EMBEDDING, Key.METADATA, Key.SCORE}

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

def test_select_all_basic_functionality():
    """Test that select_all selects all required keys on a default Search"""
    s = Search()
    codeflash_output = s.select_all(); s_all = codeflash_output # 3.64μs -> 1.65μs (121% faster)

def test_select_all_does_not_modify_original():
    """Test that select_all does not modify the original Search object"""
    s = Search()
    codeflash_output = s.select_all(); s_all = codeflash_output # 2.93μs -> 1.47μs (99.5% faster)


def test_select_all_on_search_with_empty_select():
    """Test select_all on Search with explicitly empty select"""
    s = Search(select=[])
    codeflash_output = s.select_all(); s_all = codeflash_output # 3.40μs -> 1.60μs (113% faster)

def test_select_all_on_search_with_nonstandard_keys():
    """Test select_all overrides nonstandard select keys"""
    s = Search(select=["title", "author"])
    codeflash_output = s.select_all(); s_all = codeflash_output # 2.48μs -> 1.33μs (86.1% faster)

def test_select_all_with_select_as_set():
    """Test select_all works when select is a set"""
    s = Search(select={"#document", "#score"})
    codeflash_output = s.select_all(); s_all = codeflash_output # 2.40μs -> 1.34μs (78.3% faster)

def test_select_all_with_select_as_dict():
    """Test select_all works when select is a dict"""
    s = Search(select={"keys": ["#document", "#score"]})
    codeflash_output = s.select_all(); s_all = codeflash_output # 2.30μs -> 1.30μs (76.7% faster)


def test_select_all_with_select_None():
    """Test select_all works when select is None"""
    s = Search(select=None)
    codeflash_output = s.select_all(); s_all = codeflash_output # 3.63μs -> 1.59μs (128% faster)

def test_select_all_return_type():
    """Test that select_all returns a new Search object and not self"""
    s = Search()
    codeflash_output = s.select_all(); s_all = codeflash_output # 2.94μs -> 1.44μs (104% faster)

def test_select_all_idempotency():
    """Test that calling select_all twice does not change result"""
    s = Search()
    codeflash_output = s.select_all(); s_all1 = codeflash_output # 2.76μs -> 1.40μs (97.1% faster)
    codeflash_output = s_all1.select_all(); s_all2 = codeflash_output # 1.34μs -> 568ns (137% faster)

def test_select_all_keys_are_key_objects():
    """Test that all selected keys are Key objects"""
    codeflash_output = Search().select_all(); s = codeflash_output
    for k in s.get_selected_keys():
        pass

def test_select_all_keys_are_correct():
    """Test that selected keys are exactly the expected keys"""
    codeflash_output = Search().select_all(); s = codeflash_output
    names = {k.name for k in s.get_selected_keys()}
    expected_names = {"#document", "#embedding", "#metadata", "#score"}

def test_select_all_with_malformed_select_raises():
    """Test that Search with malformed select raises TypeError"""
    with pytest.raises(TypeError):
        Search(select=123).select_all()
    with pytest.raises(TypeError):
        Search(select=object()).select_all()

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

def test_select_all_with_large_search_object():
    """Test select_all on Search with large select set (should be overridden)"""
    many_keys = [f"meta_{i}" for i in range(1000)]
    s = Search(select=many_keys)
    codeflash_output = s.select_all(); s_all = codeflash_output # 3.12μs -> 1.52μs (106% faster)

def test_select_all_performance_large_loop():
    """Test select_all performance and correctness over many Search objects"""
    searches = [Search(select=["#document", "#score"]) for _ in range(1000)]
    for s in searches:
        codeflash_output = s.select_all(); s_all = codeflash_output # 972μs -> 467μs (108% faster)

def test_select_all_does_not_leak_keys():
    """Test that select_all does not retain old keys even after many calls"""
    s = Search(select=["title", "author"])
    for _ in range(1000):
        codeflash_output = s.select_all(); s = codeflash_output # 955μs -> 462μs (107% faster)

def test_select_all_on_various_initializations():
    """Test select_all on Search objects initialized with various select types"""
    # List, set, dict, Select object, None
    selects = [
        ["#document", "#score"],
        {"#document", "#score"},
        {"keys": ["#document", "#score"]},
        Select(keys={Key.DOCUMENT, Key.SCORE}),
        None
    ]
    for sel in selects:
        s = Search(select=sel)
        codeflash_output = s.select_all(); s_all = codeflash_output


#------------------------------------------------
from chromadb.execution.expression.plan import Search

def test_Search_select_all():
    Search.select_all(Search(where=None, rank=None, limit=1, select=[]))

🔎 Concolic Coverage Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
codeflash_concolic_p_g0hne0/tmp71_sap_k/test_concolic_coverage.py::test_Search_select_all	3.13μs	1.36μs	131%✅

To edit these changes git checkout codeflash/optimize-Search.select_all-mh7jr5q1 and push.