Skip to content

Commit 73c1b82

Browse files
authored
Add native MMR support for dense search (#30)
* Add native MMR support for dense search * Add indexed_only and quantization support to search parameters
1 parent 5bfa595 commit 73c1b82

8 files changed

Lines changed: 194 additions & 21 deletions

File tree

README.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -102,6 +102,7 @@ INSERT BULK INTO COLLECTION articles VALUES [{'text': '...'}, {'text': '...'}]
102102
SEARCH articles SIMILAR TO 'query' LIMIT 10
103103
SEARCH articles SIMILAR TO 'query' LIMIT 10 WHERE year >= 2020
104104
SEARCH articles SIMILAR TO 'query' LIMIT 10 WHERE active = true
105+
SEARCH articles SIMILAR TO 'query' LIMIT 10 WITH { mmr_diversity: 0.5, mmr_candidates: 50 }
105106
SEARCH articles SIMILAR TO 'query' LIMIT 10 USING HYBRID
106107
SEARCH articles SIMILAR TO 'query' LIMIT 10 USING HYBRID FUSION 'dbsf'
107108
SEARCH articles SIMILAR TO 'query' LIMIT 10 WITH { indexed_only: true }

docs/search.md

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ SEARCH <collection_name> SIMILAR TO '<query_text>' LIMIT <n> USING HYBRID
1717
SEARCH <collection_name> SIMILAR TO '<query_text>' LIMIT <n> USING HYBRID [FUSION 'rrf|dbsf'] [DENSE MODEL '<model>'] [SPARSE MODEL '<model>'] [WHERE <filter>]
1818
SEARCH <collection_name> SIMILAR TO '<query_text>' LIMIT <n> USING SPARSE [MODEL '<sparse_model>']
1919
SEARCH <collection_name> SIMILAR TO '<query_text>' LIMIT <n> EXACT
20-
SEARCH <collection_name> SIMILAR TO '<query_text>' LIMIT <n> [USING ...] [WHERE <filter>] [RERANK] WITH { hnsw_ef: <n>, exact: true|false, acorn: true|false, indexed_only: true|false, quantization: { ignore: true|false, rescore: true|false, oversampling: <n> } }
20+
SEARCH <collection_name> SIMILAR TO '<query_text>' LIMIT <n> [USING ...] [WHERE <filter>] [RERANK] WITH { hnsw_ef: <n>, exact: true|false, acorn: true|false, indexed_only: true|false, quantization: { ignore: true|false, rescore: true|false, oversampling: <n> }, mmr_diversity: <0..1>, mmr_candidates: <n> }
2121
SEARCH <collection_name> SIMILAR TO '<query_text>' LIMIT <n> [USING ...] [WHERE <filter>] RERANK [MODEL '<reranker_model>']
2222
```
2323

@@ -55,6 +55,11 @@ Search with query-time HNSW tuning:
5555
SEARCH articles SIMILAR TO 'attention mechanism' LIMIT 10 WITH { hnsw_ef: 128 }
5656
```
5757

58+
Search with native MMR diversification:
59+
```sql
60+
SEARCH articles SIMILAR TO 'attention mechanism' LIMIT 10 WITH { mmr_diversity: 0.5, mmr_candidates: 50 }
61+
```
62+
5863
**Output:**
5964

6065
Results are displayed as a table with three columns:
@@ -102,12 +107,14 @@ Use these when you want to debug retrieval quality or tune recall without changi
102107
| `WITH { hnsw_ef: 128 }` | Increase HNSW exploration at query time |
103108
| `WITH { exact: true }` | Force exact KNN explicitly |
104109
| `WITH { acorn: true }` | Enable ACORN for filtered queries |
105-
| `WITH { indexed_only: true }` | Restrict the query to indexed segments only |
106-
| `WITH { quantization: { ... } }` | Tune quantized-search behavior at query time |
110+
| `WITH { indexed_only: true, quantization: { rescore: true } }` | Prefer indexed vectors and apply quantization controls |
111+
| `WITH { mmr_diversity: 0.5, mmr_candidates: 50 }` | Apply native MMR diversification after nearest-neighbor retrieval |
107112

108113
- `EXACT` can appear after `LIMIT` or after `RERANK`
109114
- `WITH { ... }` can appear after `WHERE` and/or `RERANK`
110-
- Supported top-level `WITH` keys are `hnsw_ef`, `exact`, `acorn`, `indexed_only`, and `quantization`
115+
- Supported top-level `WITH` keys are `hnsw_ef`, `exact`, `acorn`, `indexed_only`, `quantization`, `mmr_diversity`, and `mmr_candidates`
116+
- MMR is currently supported for dense `SEARCH` and dense `SEARCH ... GROUP BY`
117+
- MMR is not yet supported with `USING HYBRID`, `USING SPARSE`, or `RECOMMEND`
111118

112119
```sql
113120
-- Exact KNN baseline
@@ -124,6 +131,9 @@ SEARCH articles SIMILAR TO 'retrieval' LIMIT 10 WITH { indexed_only: true }
124131

125132
-- Quantized-search tuning
126133
SEARCH articles SIMILAR TO 'vector db' LIMIT 10 WITH { quantization: { ignore: true, oversampling: 2 } }
134+
135+
-- Diversify top-k results with native MMR
136+
SEARCH articles SIMILAR TO 'retrieval systems' LIMIT 10 WITH { mmr_diversity: 0.5, mmr_candidates: 50 }
127137
```
128138

129139
---

src/qql/ast_nodes.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,8 @@ class SearchWith:
2929
acorn: bool = False
3030
indexed_only: bool = False
3131
quantization: "QuantizationSearchWith | None" = None
32+
mmr_diversity: float | None = None
33+
mmr_candidates: int | None = None
3234

3335

3436
@dataclass(frozen=True)

src/qql/cli.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -70,7 +70,7 @@
7070
Optional: [yellow]WHERE[/yellow] <filter> (e.g. WHERE year > 2020 AND status = 'ok')
7171
Optional: [yellow]RERANK[/yellow] [MODEL '<model>'] rerank results with a cross-encoder
7272
Optional: [yellow]EXACT[/yellow] bypass HNSW and perform exact search
73-
Optional: [yellow]WITH[/yellow] { hnsw_ef: <int>, exact: <bool>, acorn: <bool>, indexed_only: <bool>, quantization: { ignore: <bool>, rescore: <bool>, oversampling: <n> } } search parameters
73+
Optional: [yellow]WITH[/yellow] { hnsw_ef: <int>, exact: <bool>, acorn: <bool>, indexed_only: <bool>, quantization: { ignore: <bool>, rescore: <bool>, oversampling: <n> }, mmr_diversity: <0..1>, mmr_candidates: <int> } search parameters
7474
Optional: [yellow]GROUP BY[/yellow] <field> [[yellow]GROUP_SIZE[/yellow] <n>]
7575
Group results by a payload field value (default GROUP_SIZE: 3).
7676
Field must be keyword or integer type. RERANK and GROUP BY cannot be combined.

src/qql/executor.py

Lines changed: 35 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,9 @@
2727
MatchText,
2828
MatchTextAny,
2929
MatchValue,
30+
Mmr,
3031
Modifier,
32+
NearestQuery,
3133
PayloadField,
3234
PayloadSchemaType,
3335
PointStruct,
@@ -602,6 +604,7 @@ def _execute_search(self, node: SearchStmt) -> ExecutionResult:
602604
)
603605

604606
search_params = self._build_search_params(node.with_clause)
607+
self._validate_search_mmr_usage(node)
605608

606609
# When reranking is requested, fetch more candidates so the reranker has
607610
# enough material to reorder; only `node.limit` results are returned.
@@ -712,7 +715,7 @@ def _execute_search(self, node: SearchStmt) -> ExecutionResult:
712715
query_using = self._get_dense_vector_name(node.collection)
713716
response = self._client.query_points(
714717
collection_name=node.collection,
715-
query=vector,
718+
query=self._build_dense_query(vector, node.with_clause),
716719
using=query_using,
717720
limit=fetch_limit,
718721
query_filter=qdrant_filter,
@@ -790,6 +793,8 @@ def _execute_recommend(self, node: RecommendStmt) -> ExecutionResult:
790793
)
791794

792795
search_params = self._build_search_params(node.with_clause)
796+
if self._has_mmr(node.with_clause):
797+
raise QQLRuntimeError("MMR is supported only for SEARCH statements")
793798

794799
lookup_from: LookupLocation | None = None
795800
if node.lookup_from is not None:
@@ -842,6 +847,34 @@ def _build_search_params(self, with_clause: SearchWith | None) -> SearchParams |
842847
acorn=AcornSearchParams(enable=True) if with_clause.acorn else None,
843848
)
844849

850+
def _has_mmr(self, with_clause: SearchWith | None) -> bool:
851+
return with_clause is not None and (
852+
with_clause.mmr_diversity is not None or with_clause.mmr_candidates is not None
853+
)
854+
855+
def _validate_search_mmr_usage(self, node: SearchStmt) -> None:
856+
if not self._has_mmr(node.with_clause):
857+
return
858+
if node.hybrid:
859+
raise QQLRuntimeError("MMR is not supported with USING HYBRID yet")
860+
if node.sparse_only:
861+
raise QQLRuntimeError("MMR is not supported with USING SPARSE yet")
862+
863+
def _build_dense_query(
864+
self,
865+
vector: list[float],
866+
with_clause: SearchWith | None,
867+
) -> list[float] | NearestQuery:
868+
if not self._has_mmr(with_clause):
869+
return vector
870+
return NearestQuery(
871+
nearest=vector,
872+
mmr=Mmr(
873+
diversity=with_clause.mmr_diversity,
874+
candidates_limit=with_clause.mmr_candidates,
875+
),
876+
)
877+
845878
def _parse_recommend_strategy(
846879
self, strategy: str | None
847880
) -> RecommendStrategy | None:
@@ -1029,7 +1062,7 @@ def _execute_search_groups(
10291062
response = self._client.query_points_groups(
10301063
collection_name=node.collection,
10311064
group_by=node.group_by,
1032-
query=vector,
1065+
query=self._build_dense_query(vector, node.with_clause),
10331066
using=query_using,
10341067
limit=node.limit,
10351068
group_size=node.group_size,

src/qql/parser.py

Lines changed: 28 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
QuantizationSearchWith,
2727
QuantizationConfig,
2828
QuantizationType,
29+
QuantizationSearchWith,
2930
RecommendStmt,
3031
SelectStmt,
3132
ScrollStmt,
@@ -417,6 +418,8 @@ def _parse_search(self) -> SearchStmt:
417418
acorn=with_clause.acorn,
418419
indexed_only=with_clause.indexed_only,
419420
quantization=with_clause.quantization,
421+
mmr_diversity=with_clause.mmr_diversity,
422+
mmr_candidates=with_clause.mmr_candidates,
420423
)
421424
if self._peek().kind == TokenKind.WITH:
422425
self._advance() # consume WITH
@@ -430,6 +433,12 @@ def _parse_search(self) -> SearchStmt:
430433
acorn=parsed_with.acorn or with_clause.acorn,
431434
indexed_only=parsed_with.indexed_only or with_clause.indexed_only,
432435
quantization=parsed_with.quantization or with_clause.quantization,
436+
mmr_diversity=(
437+
parsed_with.mmr_diversity
438+
if parsed_with.mmr_diversity is not None
439+
else with_clause.mmr_diversity
440+
),
441+
mmr_candidates=parsed_with.mmr_candidates or with_clause.mmr_candidates,
433442
)
434443
group_by: str | None = None
435444
group_size: int = 3
@@ -964,6 +973,8 @@ def _parse_with_clause(self) -> SearchWith:
964973
acorn: bool = False
965974
indexed_only: bool = False
966975
quantization: QuantizationSearchWith | None = None
976+
mmr_diversity: float | None = None
977+
mmr_candidates: int | None = None
967978
while self._peek().kind != TokenKind.RBRACE:
968979
key_tok = self._peek()
969980
if key_tok.kind not in (
@@ -988,10 +999,24 @@ def _parse_with_clause(self) -> SearchWith:
988999
indexed_only = self._parse_bool()
9891000
elif key == "quantization":
9901001
quantization = self._parse_quantization_search_with()
1002+
elif key == "mmr_diversity":
1003+
mmr_diversity = float(self._parse_number())
1004+
if not 0.0 <= mmr_diversity <= 1.0:
1005+
raise QQLSyntaxError(
1006+
f"mmr_diversity must be between 0 and 1, got {mmr_diversity}",
1007+
key_tok.pos,
1008+
)
1009+
elif key == "mmr_candidates":
1010+
mmr_candidates = int(self._expect(TokenKind.INTEGER).value)
1011+
if mmr_candidates <= 0:
1012+
raise QQLSyntaxError(
1013+
f"mmr_candidates must be a positive integer, got {mmr_candidates}",
1014+
key_tok.pos,
1015+
)
9911016
else:
9921017
raise QQLSyntaxError(
9931018
"Unknown WITH parameter "
994-
f"'{key}'. Expected: hnsw_ef, exact, acorn, indexed_only, quantization",
1019+
f"'{key}'. Expected: hnsw_ef, exact, acorn, indexed_only, quantization, mmr_diversity, mmr_candidates",
9951020
key_tok.pos,
9961021
)
9971022
if self._peek().kind == TokenKind.COMMA:
@@ -1007,6 +1032,8 @@ def _parse_with_clause(self) -> SearchWith:
10071032
acorn=acorn,
10081033
indexed_only=indexed_only,
10091034
quantization=quantization,
1035+
mmr_diversity=mmr_diversity,
1036+
mmr_candidates=mmr_candidates,
10101037
)
10111038

10121039
def _parse_quantization_search_with(self) -> QuantizationSearchWith:

tests/test_executor.py

Lines changed: 84 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -792,7 +792,6 @@ def test_sparse_search_forwards_search_params(self, executor, mock_client, mocke
792792
search_params = mock_client.query_points.call_args.kwargs["search_params"]
793793
assert search_params.exact is True
794794
assert search_params.indexed_only is True
795-
796795
def test_dense_search_against_hybrid_collection_uses_dense_vector_name(
797796
self, executor, mock_client, mocker
798797
):
@@ -811,6 +810,55 @@ def test_dense_search_against_hybrid_collection_uses_dense_vector_name(
811810

812811
assert mock_client.query_points.call_args.kwargs["using"] == "dense"
813812

813+
def test_dense_search_with_mmr_uses_nearest_query(self, executor, mock_client, mocker):
814+
from qdrant_client.models import NearestQuery
815+
816+
mock_client.collection_exists.return_value = True
817+
mock_response = mocker.MagicMock()
818+
mock_response.points = []
819+
mock_client.query_points.return_value = mock_response
820+
821+
node = SearchStmt(
822+
collection="notes",
823+
query_text="hello",
824+
limit=5,
825+
model=None,
826+
with_clause=SearchWith(mmr_diversity=0.4, mmr_candidates=25),
827+
)
828+
executor.execute(node)
829+
830+
query = mock_client.query_points.call_args.kwargs["query"]
831+
assert isinstance(query, NearestQuery)
832+
assert query.mmr is not None
833+
assert query.mmr.diversity == pytest.approx(0.4)
834+
assert query.mmr.candidates_limit == 25
835+
836+
def test_hybrid_search_with_mmr_raises(self, executor, mock_client):
837+
mock_client.collection_exists.return_value = True
838+
node = SearchStmt(
839+
collection="notes",
840+
query_text="hello",
841+
limit=5,
842+
model=None,
843+
hybrid=True,
844+
with_clause=SearchWith(mmr_diversity=0.5),
845+
)
846+
with pytest.raises(QQLRuntimeError, match="MMR is not supported with USING HYBRID yet"):
847+
executor.execute(node)
848+
849+
def test_sparse_search_with_mmr_raises(self, executor, mock_client):
850+
mock_client.collection_exists.return_value = True
851+
node = SearchStmt(
852+
collection="notes",
853+
query_text="hello",
854+
limit=5,
855+
model=None,
856+
sparse_only=True,
857+
with_clause=SearchWith(mmr_diversity=0.5),
858+
)
859+
with pytest.raises(QQLRuntimeError, match="MMR is not supported with USING SPARSE yet"):
860+
executor.execute(node)
861+
814862

815863
class TestRecommend:
816864
def test_recommend_calls_qdrant_query_points(self, executor, mock_client, mocker):
@@ -1026,6 +1074,17 @@ def test_recommend_forwards_indexed_only_and_quantization(self, executor, mock_c
10261074
assert search_params.quantization is not None
10271075
assert search_params.quantization.rescore is True
10281076

1077+
def test_recommend_with_mmr_raises(self, executor, mock_client):
1078+
mock_client.collection_exists.return_value = True
1079+
node = RecommendStmt(
1080+
collection="notes",
1081+
positive_ids=("a",),
1082+
limit=5,
1083+
with_clause=SearchWith(mmr_diversity=0.5),
1084+
)
1085+
with pytest.raises(QQLRuntimeError, match="MMR is supported only for SEARCH statements"):
1086+
executor.execute(node)
1087+
10291088
def test_recommend_offset_zero_passes_none(self, executor, mock_client, mocker):
10301089
mock_client.collection_exists.return_value = True
10311090
mock_response = mocker.MagicMock()
@@ -2268,12 +2327,35 @@ def test_group_by_hybrid_uses_query_points_groups(self, executor, mock_client, m
22682327
collection="articles", query_text="q", limit=3, model=None,
22692328
hybrid=True, group_by="category", group_size=2,
22702329
)
2271-
result = executor.execute(node)
2330+
executor.execute(node)
22722331
mock_client.query_points_groups.assert_called_once()
22732332
kwargs = mock_client.query_points_groups.call_args.kwargs
22742333
assert kwargs["group_by"] == "category"
22752334
assert "prefetch" in kwargs
22762335

2336+
def test_group_by_dense_with_mmr_uses_nearest_query(self, executor, mock_client, mocker):
2337+
from qdrant_client.models import NearestQuery
2338+
2339+
mock_client.collection_exists.return_value = True
2340+
mock_response = mocker.MagicMock()
2341+
mock_response.groups = []
2342+
mock_client.query_points_groups.return_value = mock_response
2343+
2344+
node = SearchStmt(
2345+
collection="articles",
2346+
query_text="ai",
2347+
limit=5,
2348+
model=None,
2349+
group_by="category",
2350+
with_clause=SearchWith(mmr_diversity=0.35, mmr_candidates=40),
2351+
)
2352+
executor.execute(node)
2353+
query = mock_client.query_points_groups.call_args.kwargs["query"]
2354+
assert isinstance(query, NearestQuery)
2355+
assert query.mmr is not None
2356+
assert query.mmr.diversity == pytest.approx(0.35)
2357+
assert query.mmr.candidates_limit == 40
2358+
22772359

22782360
class TestUpdateVector:
22792361
def test_update_vector_calls_update_vectors(self, executor, mock_client):
@@ -2288,7 +2370,6 @@ def test_update_vector_calls_update_vectors(self, executor, mock_client):
22882370

22892371
def test_update_vector_passes_correct_point_id(self, executor, mock_client):
22902372
from qql.ast_nodes import UpdateVectorStmt
2291-
from qdrant_client.models import PointVectors
22922373
mock_client.collection_exists.return_value = True
22932374
mock_client.get_collection.return_value.config.params.vectors = {} # non-dict → unnamed
22942375
node = UpdateVectorStmt(
@@ -2480,7 +2561,6 @@ def test_update_vector_unnamed_collection_sends_plain_list(self, executor, mock_
24802561
from qql.ast_nodes import UpdateVectorStmt
24812562
mock_client.collection_exists.return_value = True
24822563
# Unnamed collection: get_collection returns non-dict vectors
2483-
mock_vectors = mocker.MagicMock() if False else type("V", (), {})()
24842564
info = mock_client.get_collection.return_value
24852565
info.config.params.vectors = [None] # list → not a dict → unnamed
24862566

0 commit comments

Comments
 (0)