Skip to content

Commit fe351d4

Browse files
committed
Add moreLikeThis lookup.
1 parent f31ee2d commit fe351d4

File tree

3 files changed

+58
-15
lines changed

3 files changed

+58
-15
lines changed

django_mongodb_backend/compiler.py

Lines changed: 13 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -109,36 +109,35 @@ def _prepare_expressions_for_pipeline(self, expression, target, annotation_group
109109
replacements[sub_expr] = self._get_replace_expr(sub_expr, group, alias)
110110
return replacements, group
111111

112-
def _prepare_search_expressions_for_pipeline(self, expression, target, search_idx):
112+
def _prepare_search_expressions_for_pipeline(
113+
self, expression, target, search_idx, replacements
114+
):
113115
searches = {}
114-
replacements = {}
115116
for sub_expr in self._get_search_expressions(expression):
116-
alias = f"__search_expr.search{next(search_idx)}"
117-
replacements[sub_expr] = self._get_replace_expr(sub_expr, searches, alias)
118-
return replacements, list(searches.values())
117+
if sub_expr not in replacements:
118+
alias = f"__search_expr.search{next(search_idx)}"
119+
replacements[sub_expr] = self._get_replace_expr(sub_expr, searches, alias)
120+
return list(searches.values())
119121

120122
def _prepare_search_query_for_aggregation_pipeline(self, order_by):
121123
replacements = {}
122124
searches = []
123125
annotation_group_idx = itertools.count(start=1)
124126
for target, expr in self.query.annotation_select.items():
125-
new_replacements, expr_searches = self._prepare_search_expressions_for_pipeline(
126-
expr, target, annotation_group_idx
127+
expr_searches = self._prepare_search_expressions_for_pipeline(
128+
expr, target, annotation_group_idx, replacements
127129
)
128-
replacements.update(new_replacements)
129130
searches += expr_searches
130131

131132
for expr, _ in order_by:
132-
new_replacements, expr_searches = self._prepare_search_expressions_for_pipeline(
133-
expr, None, annotation_group_idx
133+
expr_searches = self._prepare_search_expressions_for_pipeline(
134+
expr, None, annotation_group_idx, replacements
134135
)
135-
replacements.update(new_replacements)
136136
searches += expr_searches
137137

138-
having_replacements, having_group = self._prepare_search_expressions_for_pipeline(
139-
self.having, None, annotation_group_idx
138+
having_group = self._prepare_search_expressions_for_pipeline(
139+
self.having, None, annotation_group_idx, replacements
140140
)
141-
replacements.update(having_replacements)
142141
searches += having_group
143142
return searches, replacements
144143

django_mongodb_backend/expressions/builtins.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -490,6 +490,8 @@ def as_mql(self, compiler, connection):
490490

491491

492492
class SearchMoreLikeThis(SearchExpression):
493+
search_type = "more_like_this"
494+
493495
def __init__(self, documents, score=None):
494496
self.documents = documents
495497
self.score = score

tests/queries_/test_search.py

Lines changed: 43 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
SearchGeoShape,
1212
SearchGeoWithin,
1313
SearchIn,
14+
SearchMoreLikeThis,
1415
SearchPhrase,
1516
SearchRange,
1617
SearchRegex,
@@ -27,7 +28,7 @@ def _get_collection(self, model):
2728

2829
def create_search_index(self, model, index_name, definition):
2930
collection = self._get_collection(model)
30-
idx = SearchIndexModel(definition=definition, name="test_index")
31+
idx = SearchIndexModel(definition=definition, name=index_name)
3132
collection.create_search_index(idx)
3233

3334

@@ -254,3 +255,44 @@ def test_search_geo_within(self):
254255
)
255256
)
256257
self.assertEqual(qs.first().number, 2)
258+
259+
260+
class SearchMoreLikeThisTest(TestCase, CreateIndexMixin):
261+
def setUp(self):
262+
self.create_search_index(
263+
Article,
264+
"mlt_index",
265+
{
266+
"mappings": {
267+
"dynamic": False,
268+
"fields": {"body": {"type": "string"}, "headline": {"type": "string"}},
269+
}
270+
},
271+
)
272+
self.article1 = Article.objects.create(
273+
headline="Space exploration", number=1, body="Webb telescope"
274+
)
275+
self.article2 = Article.objects.create(
276+
headline="The commodities fall",
277+
number=2,
278+
body="Commodities dropped sharply due to inflation concerns",
279+
)
280+
Article.objects.create(
281+
headline="irrelevant",
282+
number=3,
283+
body="This is a completely unrelated article about cooking",
284+
)
285+
time.sleep(1)
286+
287+
def test_search_more_like_this(self):
288+
like_docs = [
289+
{"headline": self.article1.headline, "body": self.article1.body},
290+
{"headline": self.article2.headline, "body": self.article2.body},
291+
]
292+
like_docs = [{"body": "NASA launches new satellite to explore the galaxy"}]
293+
qs = Article.objects.annotate(score=SearchMoreLikeThis(documents=like_docs)).order_by(
294+
"score"
295+
)
296+
self.assertQuerySetEqual(
297+
qs, ["space exploration", "The commodities fall"], lambda a: a.headline
298+
)

0 commit comments

Comments
 (0)