From 74bfd9e82495244345ecf7e5458b40547d96fba9 Mon Sep 17 00:00:00 2001 From: machaoxin0407 <221922045+machaoxin0407@users.noreply.github.com> Date: Thu, 1 Oct 2026 14:55:01 +0800 Subject: [PATCH 1/3] perf(search): prepare search filters once per query --- .../file_store/faiss_local_file_store.py | 14 +- .../components/file_store/local_file_store.py | 87 ++++++++- .../file_store/zvec_local_file_store.py | 12 +- scripts/benchmark_search_filter.py | 113 +++++++++++ tests/unit/test_prepared_search_filter.py | 176 ++++++++++++++++++ 5 files changed, 386 insertions(+), 16 deletions(-) create mode 100644 scripts/benchmark_search_filter.py create mode 100644 tests/unit/test_prepared_search_filter.py diff --git a/reme/components/file_store/faiss_local_file_store.py b/reme/components/file_store/faiss_local_file_store.py index 02613eab..961a9d99 100644 --- a/reme/components/file_store/faiss_local_file_store.py +++ b/reme/components/file_store/faiss_local_file_store.py @@ -3,6 +3,7 @@ import asyncio import hashlib import json +from collections.abc import Callable from contextlib import suppress from uuid import uuid4 @@ -664,6 +665,7 @@ class FaissLocalFileStore(LocalFileStore): return [] q = self._prepare(query_embedding) + matches = self._prepare_search_filter(search_filter) ntotal = index.ntotal # Small-index shortcut: when ntotal is below the brute-force threshold, @@ -683,14 +685,14 @@ class FaissLocalFileStore(LocalFileStore): if ntotal < max(1, (max(limit, ef_search)) ** 0.5) * self.hnsw_m: k = ntotal if search_filter else min(ntotal, limit + len(self._tombstones)) scores, rows = index.storage.search(q, k) - return self._collect_hits(rows[0].tolist(), scores[0].tolist(), limit, search_filter) + return self._collect_hits(rows[0].tolist(), scores[0].tolist(), limit, search_filter, matches=matches) if not search_filter: # No filter: simple over-fetch to cover tombstones. k = min(ntotal, limit + len(self._tombstones)) self._set_ef_search(index, limit) scores, rows = index.search(q, k) - return self._collect_hits(rows[0].tolist(), scores[0].tolist(), limit, search_filter) + return self._collect_hits(rows[0].tolist(), scores[0].tolist(), limit, search_filter, matches=matches) # With filter: progressively increase k until we collect enough results # or exhaust the entire index. @@ -698,7 +700,7 @@ class FaissLocalFileStore(LocalFileStore): while True: self._set_ef_search(index, limit) scores, rows = index.search(q, k) - results = self._collect_hits(rows[0].tolist(), scores[0].tolist(), limit, search_filter) + results = self._collect_hits(rows[0].tolist(), scores[0].tolist(), limit, search_filter, matches=matches) if len(results) >= limit or k >= ntotal: return results k = min(ntotal, k * 2) @@ -709,15 +711,19 @@ class FaissLocalFileStore(LocalFileStore): scores: list[float], limit: int, search_filter: dict | None = None, + *, + matches: Callable[[FileChunk], bool] | None = None, ) -> list[FileChunk]: """Map raw FAISS rows back to chunks, skipping tombstones and stale ids.""" + if matches is None: + matches = self._prepare_search_filter(search_filter) results: list[FileChunk] = [] for raw_row, score in zip(rows, scores): row = int(raw_row) if row < 0 or row in self._tombstones or row >= len(self._id_map): continue chunk = self.file_chunks.get(self._id_map[row]) - if chunk is None or not self._matches_search_filter(chunk, search_filter): + if chunk is None or not matches(chunk): continue results.append(chunk.model_copy(update={"scores": {"vector": float(score), "score": float(score)}})) if len(results) >= limit: diff --git a/reme/components/file_store/local_file_store.py b/reme/components/file_store/local_file_store.py index 243917f8..33509100 100644 --- a/reme/components/file_store/local_file_store.py +++ b/reme/components/file_store/local_file_store.py @@ -6,7 +6,7 @@ import datetime import heapq import json import time -from collections.abc import Iterable +from collections.abc import Callable, Iterable from contextlib import suppress import numpy as np @@ -941,6 +941,7 @@ class LocalFileStore(BaseFileStore): if query_embedding is None: return [] + matches = self._prepare_search_filter(search_filter) top: list[tuple[float, int, FileChunk]] = [] candidates: list[FileChunk] = [] embeddings: list[np.ndarray] = [] @@ -964,10 +965,7 @@ class LocalFileStore(BaseFileStore): embeddings.clear() for candidate in self.file_chunks.values(): - if not self._embedding_dim_matches(candidate.embedding) or not self._matches_search_filter( - candidate, - search_filter, - ): + if not self._embedding_dim_matches(candidate.embedding) or not matches(candidate): continue candidates.append(candidate) embeddings.append(candidate.embedding) @@ -989,10 +987,9 @@ class LocalFileStore(BaseFileStore): if not query: return [] + matches = self._prepare_search_filter(search_filter) if search_filter: - eligible_ids = { - chunk.id for chunk in self.file_chunks.values() if self._matches_search_filter(chunk, search_filter) - } + eligible_ids = {chunk.id for chunk in self.file_chunks.values() if matches(chunk)} if not eligible_ids: return [] try: @@ -1011,7 +1008,7 @@ class LocalFileStore(BaseFileStore): results = [] for doc_id, score in doc_id_score_dict.items(): chunk = self.file_chunks.get(doc_id) - if chunk and self._matches_search_filter(chunk, search_filter): + if chunk and matches(chunk): results.append(chunk.model_copy(update={"scores": {"keyword": score, "score": score}})) if len(results) >= limit: break @@ -1020,6 +1017,78 @@ class LocalFileStore(BaseFileStore): # -- extensions ----------------------------------------------------------- + @classmethod + def _prepare_search_filter(cls, search_filter: dict | None) -> Callable[[FileChunk], bool]: + """Prepare query-local membership tests; retain the dict-based compatibility path.""" + + def legacy(chunk): + return cls._matches_search_filter(chunk, search_filter) + + # Subclasses may customize matching without implementing this new helper. + if ( + cls._matches_search_filter.__func__ is not LocalFileStore._matches_search_filter.__func__ + or cls._value_matches.__func__ is not LocalFileStore._value_matches.__func__ + ): + return legacy + if not search_filter: + return lambda chunk: True + try: + exact_paths = set() + has_exact_paths = False + for key in ("path", "paths"): + if key in search_filter: + has_exact_paths = True + exact_paths.update(cls._as_filter_values(search_filter[key])) + prefixes = tuple( + str(value) + for key in ("path_prefix", "path_prefixes", "prefix", "prefixes") + if key in search_filter + for value in cls._as_filter_values(search_filter[key]) + ) + start_date = search_filter.get("start_date") + end_date = search_filter.get("end_date") + strict_date = bool(search_filter.get("strict_date_filter", False)) + metadata = dict(search_filter.get("metadata") or {}) + reserved = { + "path", + "paths", + "path_prefix", + "path_prefixes", + "prefix", + "prefixes", + "metadata", + "start_date", + "end_date", + "strict_date_filter", + } + metadata.update((key, value) for key, value in search_filter.items() if key not in reserved) + conditions = [] + for key, value in metadata.items(): + collection = isinstance(value, (list, tuple, set, frozenset)) + conditions.append((key, frozenset(value) if collection else value, collection)) + except (TypeError, ValueError): + # Malformed late-stage conditions must still short-circuit as before. + return legacy + + def matches(chunk: FileChunk) -> bool: + if has_exact_paths and chunk.path not in exact_paths: + return False + if prefixes and not chunk.path.startswith(prefixes): + return False + if start_date or end_date: + path_date = cls._extract_date_from_path(chunk.path) + if not path_date: + if strict_date: + return False + elif (start_date and path_date < start_date) or (end_date and path_date > end_date): + return False + return all( + chunk.metadata.get(key) in expected if collection else chunk.metadata.get(key) == expected + for key, expected, collection in conditions + ) + + return matches + @staticmethod def _as_filter_values(value) -> set: if isinstance(value, (list, tuple, set, frozenset)): diff --git a/reme/components/file_store/zvec_local_file_store.py b/reme/components/file_store/zvec_local_file_store.py index 58a0300a..6e5151e5 100644 --- a/reme/components/file_store/zvec_local_file_store.py +++ b/reme/components/file_store/zvec_local_file_store.py @@ -4,6 +4,7 @@ import hashlib import json import shutil import time +from collections.abc import Callable from uuid import uuid4 import aiofiles @@ -455,18 +456,19 @@ class ZvecLocalFileStore(LocalFileStore): return [] vector = np.asarray(query_embedding, dtype=np.float32).tolist() + matches = self._prepare_search_filter(search_filter) ntotal = len(self._indexed_ids) if not search_filter: hits = self._query_collection(collection, vector, min(limit, ntotal)) - return self._collect_hits(hits, limit, search_filter) + return self._collect_hits(hits, limit, search_filter, matches=matches) # With a post-filter: progressively increase k until we collect enough # results or exhaust the reachable index. k = min(ntotal, 3 * limit) while True: hits = self._query_collection(collection, vector, k) - results = self._collect_hits(hits, limit, search_filter) + results = self._collect_hits(hits, limit, search_filter, matches=matches) if len(results) >= limit or k >= ntotal: return results k = min(ntotal, k * 2) @@ -484,12 +486,16 @@ class ZvecLocalFileStore(LocalFileStore): hits: list[tuple[str, float]], limit: int, search_filter: dict | None = None, + *, + matches: Callable[[FileChunk], bool] | None = None, ) -> list[FileChunk]: """Map zvec hits back to chunks, skipping stale ids and filtered-out chunks.""" + if matches is None: + matches = self._prepare_search_filter(search_filter) results: list[FileChunk] = [] for chunk_id, score in hits: chunk = self.file_chunks.get(chunk_id) - if chunk is None or not self._matches_search_filter(chunk, search_filter): + if chunk is None or not matches(chunk): continue results.append(chunk.model_copy(update={"scores": {"vector": score, "score": score}})) if len(results) >= limit: diff --git a/scripts/benchmark_search_filter.py b/scripts/benchmark_search_filter.py new file mode 100644 index 00000000..f9195be7 --- /dev/null +++ b/scripts/benchmark_search_filter.py @@ -0,0 +1,113 @@ +"""Compare prepared filtering with the unchanged dict matcher on real searches. + +Run: python -m scripts.benchmark_search_filter --repeats 7 +Setup is excluded; deterministic chunks and fake query embeddings need no files/network. +""" + +# pylint: disable=protected-access + +import argparse +import asyncio +import json +import platform +import statistics +import time +from types import SimpleNamespace +from unittest.mock import AsyncMock, patch + +import numpy as np + +from reme.components.file_store import LocalFileStore +from reme.components.keyword_index import BM25Index +from reme.components.tokenizer import RegexTokenizer +from reme.schema import FileChunk + + +async def measure(operation, repeats): + """Warm up and report per-call latency in milliseconds.""" + for _ in range(2): + await operation() + samples = [] + for _ in range(repeats): + start = time.perf_counter() + await operation() + samples.append((time.perf_counter() - start) * 1000) + return {"median_ms": statistics.median(samples), "p95_ms": float(np.percentile(samples, 95))} + + +async def benchmark(size, repeats): + """Measure predicate traversal, keyword_search and vector_search separately.""" + store = LocalFileStore(embedding_store="") + index = BM25Index() + index.tokenizer = RegexTokenizer(filter_stopwords=False) + store.keyword_index = index + store.embedding_store = SimpleNamespace(dimensions=2, is_healthy=True) + store._get_query_embedding = AsyncMock(return_value=np.array([1.0, 0.0])) + for i in range(size): + store.file_chunks[str(i)] = FileChunk( + id=str(i), + path=f"daily/{i % 2000}.md", + text="alpha beta", + embedding=np.array([1.0, 0.0]), + ) + await index.add_docs({key: chunk.text for key, chunk in store.file_chunks.items()}) + results = [] + for path_count in (1, 100, 1000): + config = {"paths": [f"daily/{i}.md" for i in range(path_count)]} + + async def predicate(config=config): + matches = store._prepare_search_filter(config) + return [chunk.id for chunk in store.file_chunks.values() if matches(chunk)] + + async def keyword(config=config): + return [(c.id, c.scores) for c in await store.keyword_search("alpha", 10, config)] + + async def vector(config=config): + return [(c.id, c.scores) for c in await store.vector_search("alpha", 10, config)] + + def legacy(search_filter): + return lambda chunk: store._matches_search_filter(chunk, search_filter) + + for name, operation in (("predicate", predicate), ("keyword_search", keyword), ("vector_search", vector)): + expected = await operation() + prepared = await measure(operation, repeats) + with patch.object(store, "_prepare_search_filter", side_effect=legacy): + assert await operation() == expected + baseline = await measure(operation, repeats) + results.append( + { + "chunks": size, + "paths": path_count, + "operation": name, + "baseline": baseline, + "prepared": prepared, + }, + ) + return results + + +async def main(): + """Print reproducible environment and results as JSON.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--sizes", nargs="+", type=int, default=[20000, 100000]) + parser.add_argument("--repeats", type=int, default=7) + args = parser.parse_args() + rows = [] + for size in args.sizes: + rows.extend(await benchmark(size, args.repeats)) + print( + json.dumps( + { + "python": platform.python_version(), + "platform": platform.platform(), + "numpy": np.__version__, + "repeats": args.repeats, + "results": rows, + }, + indent=2, + ), + ) + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/tests/unit/test_prepared_search_filter.py b/tests/unit/test_prepared_search_filter.py new file mode 100644 index 00000000..d213c7ab --- /dev/null +++ b/tests/unit/test_prepared_search_filter.py @@ -0,0 +1,176 @@ +"""Prepared filters preserve the dict matcher and all search entrypoints.""" + +# pylint: disable=protected-access,missing-function-docstring + +import copy +import random +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock, patch + +import numpy as np +import pytest + +from reme.components.file_store import FaissLocalFileStore, LocalFileStore, ZvecLocalFileStore +from reme.components.keyword_index import BM25Index +from reme.components.tokenizer import RegexTokenizer +from reme.schema import FileChunk + + +def test_differential_filter_combinations(): + rng = random.Random(714) + paths = ["daily/2026-01-01/n.md", "daily/2026-02-01.md", "digest/topic.md", "daily/invalid.md"] + chunks = [ + FileChunk(id=str(i), path=rng.choice(paths), metadata={"group": rng.choice([None, "a", "b"]), "n": i % 3}) + for i in range(80) + ] + path_filters = [{}, {"paths": []}, {"path": paths[0]}, {"path": paths[0], "paths": paths[1:]}] + prefix_filters = [{}, {"prefixes": []}, {"prefix": "daily/"}, {"path_prefix": "digest", "prefixes": ["daily/"]}] + date_filters = [ + {}, + {"strict_date_filter": True}, + {"start_date": "2026-01-15"}, + { + "end_date": "2026-01-31", + "strict_date_filter": True, + }, + ] + metadata_filters = [ + {}, + {"group": ["a", "b"]}, + {"metadata": {"group": "a"}, "group": "b"}, + { + "metadata": {"n": (1, 2), "absent": None}, + }, + ] + for _ in range(250): + config = { + **rng.choice(path_filters), + **rng.choice(prefix_filters), + **rng.choice(date_filters), + **rng.choice(metadata_filters), + } + original = copy.deepcopy(config) + prepared = LocalFileStore._prepare_search_filter(config) + assert [prepared(chunk) for chunk in chunks] == [ + LocalFileStore._matches_search_filter(chunk, config) for chunk in chunks + ] + assert config == original + assert LocalFileStore._prepare_search_filter(None)(chunks[0]) + + +def test_preparation_is_query_local_and_preserves_subclass_hooks(): + candidate = FileChunk(path="one") + config = {"paths": ["one"]} + prepared = LocalFileStore._prepare_search_filter(config) + config["paths"][:] = ["two"] + assert prepared(candidate) + assert not LocalFileStore._prepare_search_filter(config)(candidate) + + class CustomStore(LocalFileStore): + """Existing subclass hook still governs the new hot path.""" + + @classmethod + def _matches_search_filter(cls, chunk, search_filter): + return chunk.path == "special" + + assert not CustomStore._prepare_search_filter({})(candidate) + assert CustomStore._prepare_search_filter({})(FileChunk(path="special")) + + +@pytest.mark.parametrize( + "config", + [ + {"paths": [], "metadata": {"bad": [[1]]}}, + {"paths": ["one"], "metadata": {"bad": [[1]]}}, + {"prefix": "else", "metadata": 123}, + {"paths": [[1]]}, + ], +) +def test_malformed_conditions_keep_short_circuit_behavior(config): + candidate = FileChunk(path="one") + prepared = LocalFileStore._prepare_search_filter(config) + try: + expected = LocalFileStore._matches_search_filter(candidate, config) + except (TypeError, ValueError) as error: + with pytest.raises(type(error)): + prepared(candidate) + else: + assert prepared(candidate) == expected + + +def make_store(cls=LocalFileStore): + # Native ANN boundaries are faked below; execute the real backend search methods. + store = object.__new__(cls) + LocalFileStore.__init__(store, embedding_store="") + store.file_chunks = { + str(i): FileChunk(id=str(i), path=f"daily/{i}.md", text="alpha beta", embedding=np.array([1.0, 0.0])) + for i in range(30) + } + store.embedding_store = SimpleNamespace(dimensions=2, is_healthy=True) + store._get_query_embedding = AsyncMock(return_value=np.array([1.0, 0.0])) + return store + + +@pytest.mark.asyncio +@pytest.mark.parametrize("fallback", [False, True]) +async def test_keyword_prepares_once_and_preserves_fallback(fallback): + store = make_store() + index = BM25Index() + index.tokenizer = RegexTokenizer(filter_stopwords=False) + await index.add_docs({key: chunk.text for key, chunk in store.file_chunks.items()}) + store.keyword_index = index + if fallback: + index.retrieve_filtered = AsyncMock(side_effect=NotImplementedError) + config = {"paths": ["daily/28.md", "daily/29.md"]} + with patch.object(store, "_prepare_search_filter", wraps=store._prepare_search_filter) as prepare: + result = await store.keyword_search("alpha", 2, config) + assert prepare.call_count == 1 + with patch.object( + store, + "_prepare_search_filter", + side_effect=lambda filt: lambda c: store._matches_search_filter(c, filt), + ): + original = await store.keyword_search("alpha", 2, config) + assert [(c.id, c.scores) for c in result] == [(c.id, c.scores) for c in original] + assert {c.path for c in result} == set(config["paths"]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("backend", [LocalFileStore, FaissLocalFileStore, ZvecLocalFileStore]) +async def test_vector_prepares_once_including_progressive_recall(backend): + store = make_store(backend) + searches = [] + if backend is FaissLocalFileStore: + + def search(_query, count): + searches.append(count) + return np.ones((1, count)), np.arange(count).reshape(1, -1) + + store._faiss_index = SimpleNamespace(ntotal=30, search=search, hnsw=SimpleNamespace(efSearch=0)) + store._prepare = Mock(side_effect=lambda value: value.reshape(1, -1)) + store.hnsw_m = 1 + store._id_map = list(store.file_chunks) + store._tombstones = set() + elif backend is ZvecLocalFileStore: + + def query(_collection, _vector, count): + searches.append(count) + return [(str(i), 1.0) for i in range(count)] + + store._collection = object() + store._indexed_ids = set(store.file_chunks) + store._query_collection = Mock(side_effect=query) + config = {"paths": ["daily/29.md"]} + with patch.object(store, "_prepare_search_filter", wraps=store._prepare_search_filter) as prepare: + result = await store.vector_search("alpha", 1, config) + assert prepare.call_count == 1 + if backend is not LocalFileStore: + assert searches == [3, 6, 12, 24, 30] + with patch.object( + store, + "_prepare_search_filter", + side_effect=lambda filt: lambda c: store._matches_search_filter(c, filt), + ): + original = await store.vector_search("alpha", 1, config) + assert [(c.id, c.scores) for c in result] == [(c.id, c.scores) for c in original] + assert [c.id for c in result] == ["29"] From c66b38b5d0778aa734b7010d3cf865cd85925e64 Mon Sep 17 00:00:00 2001 From: machaoxin0407 <221922045+machaoxin0407@users.noreply.github.com> Date: Thu, 1 Oct 2026 14:57:54 +0800 Subject: [PATCH 2/3] test(search): bind benchmark cases without mutable defaults --- scripts/benchmark_search_filter.py | 30 ++++++++++++++++-------------- 1 file changed, 16 insertions(+), 14 deletions(-) diff --git a/scripts/benchmark_search_filter.py b/scripts/benchmark_search_filter.py index f9195be7..d4c25bd4 100644 --- a/scripts/benchmark_search_filter.py +++ b/scripts/benchmark_search_filter.py @@ -12,6 +12,7 @@ import json import platform import statistics import time +from functools import partial from types import SimpleNamespace from unittest.mock import AsyncMock, patch @@ -51,24 +52,25 @@ async def benchmark(size, repeats): embedding=np.array([1.0, 0.0]), ) await index.add_docs({key: chunk.text for key, chunk in store.file_chunks.items()}) + + async def predicate(config): + matches = store._prepare_search_filter(config) + return [chunk.id for chunk in store.file_chunks.values() if matches(chunk)] + + async def keyword(config): + return [(c.id, c.scores) for c in await store.keyword_search("alpha", 10, config)] + + async def vector(config): + return [(c.id, c.scores) for c in await store.vector_search("alpha", 10, config)] + + def legacy(search_filter): + return lambda chunk: store._matches_search_filter(chunk, search_filter) + results = [] for path_count in (1, 100, 1000): config = {"paths": [f"daily/{i}.md" for i in range(path_count)]} - - async def predicate(config=config): - matches = store._prepare_search_filter(config) - return [chunk.id for chunk in store.file_chunks.values() if matches(chunk)] - - async def keyword(config=config): - return [(c.id, c.scores) for c in await store.keyword_search("alpha", 10, config)] - - async def vector(config=config): - return [(c.id, c.scores) for c in await store.vector_search("alpha", 10, config)] - - def legacy(search_filter): - return lambda chunk: store._matches_search_filter(chunk, search_filter) - for name, operation in (("predicate", predicate), ("keyword_search", keyword), ("vector_search", vector)): + operation = partial(operation, config) expected = await operation() prepared = await measure(operation, repeats) with patch.object(store, "_prepare_search_filter", side_effect=legacy): From b3cdc52b44dacf8aee8723f92521c166065e4d8f Mon Sep 17 00:00:00 2001 From: machaoxin0407 <221922045+machaoxin0407@users.noreply.github.com> Date: Thu, 1 Oct 2026 14:59:27 +0800 Subject: [PATCH 3/3] test(search): pass benchmark filters explicitly --- scripts/benchmark_search_filter.py | 16 +++++++--------- 1 file changed, 7 insertions(+), 9 deletions(-) diff --git a/scripts/benchmark_search_filter.py b/scripts/benchmark_search_filter.py index d4c25bd4..58b2f1f2 100644 --- a/scripts/benchmark_search_filter.py +++ b/scripts/benchmark_search_filter.py @@ -12,7 +12,6 @@ import json import platform import statistics import time -from functools import partial from types import SimpleNamespace from unittest.mock import AsyncMock, patch @@ -24,14 +23,14 @@ from reme.components.tokenizer import RegexTokenizer from reme.schema import FileChunk -async def measure(operation, repeats): +async def measure(operation, config, repeats): """Warm up and report per-call latency in milliseconds.""" for _ in range(2): - await operation() + await operation(config) samples = [] for _ in range(repeats): start = time.perf_counter() - await operation() + await operation(config) samples.append((time.perf_counter() - start) * 1000) return {"median_ms": statistics.median(samples), "p95_ms": float(np.percentile(samples, 95))} @@ -70,12 +69,11 @@ async def benchmark(size, repeats): for path_count in (1, 100, 1000): config = {"paths": [f"daily/{i}.md" for i in range(path_count)]} for name, operation in (("predicate", predicate), ("keyword_search", keyword), ("vector_search", vector)): - operation = partial(operation, config) - expected = await operation() - prepared = await measure(operation, repeats) + expected = await operation(config) + prepared = await measure(operation, config, repeats) with patch.object(store, "_prepare_search_filter", side_effect=legacy): - assert await operation() == expected - baseline = await measure(operation, repeats) + assert await operation(config) == expected + baseline = await measure(operation, config, repeats) results.append( { "chunks": size,