Python: In-memory vector search copies every match before paging #8543

Motivation & Context

Python InMemoryCollection.search() deep-copies every record that passes filtering and the score threshold, then sorts and pages the copies. A request for a small page therefore allocates copies of the vectors and nested data belonging to all matching records. This adds avoidable memory and CPU cost to local vector retrieval.

I am working on a small change that defers defensive copying until the result page has been selected. Sorting, score thresholds, total-count metadata, and detached returned records retain their existing behavior.

Reproduction

Environment: Windows x64, Python 3.12.13, agent-framework-core 1.19.0, main commit 723256961e8e46980b869ee2c75ef05f16644921.

import asyncio
import tracemalloc

from agent_framework import InMemoryCollection, VectorStoreCollectionDefinition, VectorStoreField


async def main():
    definition = VectorStoreCollectionDefinition(
        [
            VectorStoreField("key", name="id"),
            VectorStoreField("data", name="text"),
            VectorStoreField("vector", name="vector", dimensions=768, distance_function="dot_prod"),
        ],
        collection_name="search-page-benchmark",
    )
    collection = InMemoryCollection(dict, definition=definition)
    await collection.ensure_collection_exists()
    await collection.upsert(
        [
            {"id": str(i), "text": f"Document {i}", "vector": [float(i + 1)] * 768}
            for i in range(2000)
        ],
        generate_vectors=False,
    )
    query = [1.0] * 768
    # Warm up the public search and result conversion paths.
    warmup = await collection.search(vector=query, top=3, skip=4, include_vectors=True)
    _ = [row async for row in warmup]
    tracemalloc.start()
    result = await collection.search(vector=query, top=3, skip=4, include_vectors=True)
    rows = [row async for row in result]
    print([row["record"]["id"] for row in rows])
    print(result.metadata)
    print("Peak bytes:", tracemalloc.get_traced_memory()[1])
    tracemalloc.stop()


asyncio.run(main())

Observed results

A benchmark using the same 2,000 records, 768-dimensional vectors, top=3, and skip=4 measured 14,696,088 peak allocated bytes on current main. Deferring the copies to the selected page reduced that to 463,508 bytes. All 2,000 matches remain represented in total-count metadata, and the selected IDs and scores are identical.

Separately timed searches, with allocation tracing disabled, had five-run medians of 3.098 seconds before the change and 2.046 seconds after it on this machine. These are local measurements for this workload.

Expected behavior

Defensive copying should scale with the selected page size. Returned mutable data should remain independent of stored records.