server: optimize similar search to max 3 tag removals

2 files changed, 22 insertions, 2 deletions

diff --git a/server/szurubooru/func/similar.py b/server/szurubooru/func/similar.py
index 3955c03..8247de8 100644
--- a/server/szurubooru/func/similar.py
+++ b/server/szurubooru/func/similar.py
@@ -6,16 +6,20 @@ from szurubooru import model, search
 _search_executor_config = search.configs.PostSearchConfig()
 _search_executor = search.Executor(_search_executor_config)
 
+_max_removals = 3
+
 
 def find_similar_posts(source_post: model.Post, limit: int) -> List[model.Post]:
     results = []
     queue = Queue()  # contains lists of tags to search
     queue.put(source_post.tags)
+    source_tag_count = len(source_post.tags)
 
     while not queue.empty():
         # put follow-up searches on the queue
         last_tags = queue.get()
-        if len(last_tags) > 1:
+        tag_count = len(last_tags)
+        if tag_count > 1 and tag_count > source_tag_count - _max_removals:
             for removed_tag in last_tags:
                 next_search = list(filter(lambda t: t != removed_tag, last_tags))
                 queue.put(next_search)
@@ -27,7 +31,7 @@ def find_similar_posts(source_post: model.Post, limit: int) -> List[model.Post]:
             query += ' -id:%d' % r.post_id
 
         # execute
-        _, posts = _search_executor.execute(query, 0, limit)
+        _, posts = _search_executor.execute(query, 0, limit - len(results))
 
         # update results
         for p in posts:
diff --git a/server/szurubooru/tests/func/test_similar.py b/server/szurubooru/tests/func/test_similar.py
index 1a147db..9200ff1 100644
--- a/server/szurubooru/tests/func/test_similar.py
+++ b/server/szurubooru/tests/func/test_similar.py
@@ -41,3 +41,19 @@ def test_find_similar_posts(post_factory, tag_factory, verify_posts):
 
     results = similar.find_similar_posts(postAC, 10)
     verify_posts(results, [postABC, postBC, postAB, postA])
+
+
+def test_find_similar_posts_with_limit(post_factory, tag_factory, verify_posts):
+    tagA = tag_factory(names=["a"])
+    tagB = tag_factory(names=["b"])
+    tagC = tag_factory(names=["c"])
+    tagD = tag_factory(names=["d"])
+    tagE = tag_factory(names=["e"])
+    postA = post_factory(id=111, tags=[tagA])
+    postAB = post_factory(id=112, tags=[tagA, tagB])
+    postABCDE = post_factory(id=113, tags=[tagA, tagB, tagC, tagD, tagE])
+    db.session.add_all([tagA, tagB, tagC, tagD, tagE, postA, postAB, postABCDE])
+    db.session.flush()
+
+    results = similar.find_similar_posts(postABCDE, 10)
+    verify_posts(results, [postAB])