Filtered ANN Search With Metadata Predicatemedium

Filtered ANN Search With Metadata Predicate

medium RAG & Retrieval 15 min

NumPy allowed · plain Python earns 10 bonus points

A query retrieves context

Embedding similarity ranks the documents before generation.

Open lesson →

01query

02embed

03retrieve

04context

What is RAG?→ embed →top context

retrieval.md44%

rag-systems.md33%

agents.md24%

Cosine-like similarity scores ranked before context assembly

Implement pre-filtered ANN search: real vector search almost always needs to combine similarity ranking with a metadata constraint (e.g. 'only in-stock items').

Task

Implement filtered_topk(query, items, k, category). items is a list of {"id": str, "vector": list[float], "category": str} dicts. First filter to items whose category matches exactly, then rank the survivors by cosine similarity to query descending, and return the top-k ids. Ties broken by id ascending. If fewer than k items survive the filter, return all of them.

Function Signature

filtered_topk(query: list[float], items: list[dict], k: int, category: str) -> list[str]

Examples

Example 1: basic filter+rank
Input: query = [1, 0]
items = [{'id': 'a', 'vector': [1, 0], 'category': 'shoes'}, {'id': 'b', 'vector': [0, 1], 'category': 'hats'}]
k = 2
category = 'shoes'
Output: ['a']

Example 2: fewer survivors than k
Input: query = [1, 0]
items = [{'id': 'a', 'vector': [1, 0], 'category': 'x'}]
k = 5
category = 'x'
Output: ['a']

Example 3: no survivors
Input: query = [1, 0]
items = [{'id': 'a', 'vector': [1, 0], 'category': 'x'}]
k = 5
category = 'y'
Output: []

Constraints

1 <= k
all vectors share query's dimensionality and are non-zero

Python3Saved ✓

def filtered_topk(query: list[float], items: list[dict], k: int, category: str) -> list[str]:

# Your implementation here

pass

Case 1: basic filter+rank

Input: query = [1, 0]
items = [{'id': 'a', 'vector': [1, 0], 'category': 'shoes'}, {'id': 'b', 'vector': [0, 1], 'category': 'hats'}]
k = 2
category = 'shoes'

Expected: ['a']

Case 2: fewer survivors than k

Input: query = [1, 0]
items = [{'id': 'a', 'vector': [1, 0], 'category': 'x'}]
k = 5
category = 'x'

Expected: ['a']

Case 3: no survivors

Input: query = [1, 0]
items = [{'id': 'a', 'vector': [1, 0], 'category': 'x'}]
k = 5
category = 'y'

Expected: []