server: implement similar posts as a single query

2 files changed, 20 insertions, 69 deletions

diff --git a/server/szurubooru/func/similar.py b/server/szurubooru/func/similar.py
index 2e100ab..1a84e79 100644
--- a/server/szurubooru/func/similar.py
+++ b/server/szurubooru/func/similar.py
@@ -1,8 +1,8 @@
-from math import ceil
-from queue import Queue
 from typing import List
 
-from szurubooru import model, search
+import sqlalchemy as sa
+
+from szurubooru import db, model, search
 
 _search_executor_config = search.configs.PostSearchConfig()
 _search_executor = search.Executor(_search_executor_config)
@@ -11,33 +11,20 @@ _search_executor = search.Executor(_search_executor_config)
 def find_similar_posts(
     source_post: model.Post, limit: int, query_text: str = ''
 ) -> List[model.Post]:
-    results = []
-    # Sort tags in order of increasing post count, i.e. least to most popular
-    # This will help yield results quicker
-    source_tags = sorted(source_post.tags, key=lambda t: t.post_count)
-    source_tag_count = len(source_tags)
-    max_removals = source_tag_count - 1
-
-    tags = source_tags
-    for x in range(max_removals + 1):
-        # prepare the current search, remove known results
-        query = query_text + ' ' + ' '.join([t.first_name for t in tags])
-        query += ' -id:%d' % source_post.post_id
-        for r in results:
-            query += ' -id:%d' % r.post_id
-
-        # execute
-        _, posts = _search_executor.execute(query, 0, limit - len(results))
-
-        # update results
-        for p in posts:
-            results.append(p)
-            if len(results) >= limit:
-                break
-
-        # remove the least popular tag
-        if len(tags) <= 1:
-            break
-        tags = tags[1:]
-
-    return results
+    post_alias = sa.orm.aliased(model.Post)
+    pt_alias = sa.orm.aliased(model.PostTag)
+    result = (
+        db.session.query(post_alias)
+        .join(pt_alias, pt_alias.post_id == post_alias.post_id)
+        .filter(
+            sa.sql.or_(
+                pt_alias.tag_id == tag.tag_id for tag in source_post.tags
+            )
+        )
+        .filter(pt_alias.post_id != source_post.post_id)
+        .group_by(post_alias.post_id)
+        .order_by(sa.func.count(pt_alias.tag_id).desc())
+        .order_by(post_alias.post_id.desc())
+        .limit(limit)
+    )
+    return result
diff --git a/server/szurubooru/tests/func/test_similar.py b/server/szurubooru/tests/func/test_similar.py
index 486f67b..5346cb0 100644
--- a/server/szurubooru/tests/func/test_similar.py
+++ b/server/szurubooru/tests/func/test_similar.py
@@ -13,7 +13,6 @@ def verify_posts():
     return verify
 
 
-# I'd like my behavior to be like this
 def test_find_similar_posts(post_factory, tag_factory, verify_posts):
     tagA = tag_factory(names=["a"])
     tagB = tag_factory(names=["b"])
@@ -45,38 +44,6 @@ def test_find_similar_posts(post_factory, tag_factory, verify_posts):
     verify_posts(results, [postABC, postBC, postAB, postA])
 
 
-# but it's actually like this for performance reasons
-def test_find_similar_posts_naive(post_factory, tag_factory, verify_posts):
-    tagA = tag_factory(names=["a"])  # count=4
-    tagB = tag_factory(names=["b"])  # count=3
-    tagC = tag_factory(names=["c"])  # count=3
-    postA = post_factory(id=1, tags=[tagA])
-    postAB = post_factory(id=2, tags=[tagA, tagB])
-    postAC = post_factory(id=3, tags=[tagA, tagC])
-    postABC = post_factory(id=4, tags=[tagA, tagB, tagC])
-    postBC = post_factory(id=5, tags=[tagB, tagC])
-    db.session.add_all([tagA, tagB, tagC, postA, postAB, postAC, postABC, postBC])
-    db.session.flush()
-
-    results = similar.find_similar_posts(postBC, 10)
-    verify_posts(results, [postABC, postAC])
-
-    results = similar.find_similar_posts(postBC, 2)
-    verify_posts(results, [postABC, postAC])
-
-    results = similar.find_similar_posts(postABC, 10)
-    verify_posts(results, [postAC, postAB, postA])
-
-    results = similar.find_similar_posts(postA, 10)
-    verify_posts(results, [postABC, postAC, postAB])  # sorted by id
-
-    results = similar.find_similar_posts(postAB, 10)
-    verify_posts(results, [postABC, postAC, postA])
-
-    results = similar.find_similar_posts(postAC, 10)
-    verify_posts(results, [postABC, 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"])
@@ -90,7 +57,4 @@ def test_find_similar_posts_with_limit(post_factory, tag_factory, verify_posts):
     db.session.flush()
 
     results = similar.find_similar_posts(postABCDE, 10)
-    # I'd like it to be like this:
-    # verify_posts(results, [postAB])
-    # but it's like this for performance reasons:
     verify_posts(results, [postAB, postA])