[#477] User search improvements: tsquery search with field weights, friends & followers boosting.

lib/pleroma/user.ex

@@ -35,7 +35,7 @@ defmodule Pleroma.User do
     field(:avatar, :map)
     field(:local, :boolean, default: true)
     field(:follower_address, :string)
-    field(:search_distance, :float, virtual: true)
+    field(:search_rank, :float, virtual: true)
     field(:tags, {:array, :string}, default: [])
     field(:last_refreshed_at, :naive_datetime)
     has_many(:notifications, Notification)
@@ -511,6 +511,12 @@ defmodule Pleroma.User do
     {:ok, Repo.all(q)}
   end
 
+  def get_followers_ids(user, page \\ nil) do
+    q = get_followers_query(user, page)
+
+    Repo.all(from(u in q, select: u.id))
+  end
+
   def get_friends_query(%User{id: id, following: following}, nil) do
     from(
       u in User,
@@ -535,6 +541,12 @@ defmodule Pleroma.User do
     {:ok, Repo.all(q)}
   end
 
+  def get_friends_ids(user, page \\ nil) do
+    q = get_friends_query(user, page)
+
+    Repo.all(from(u in q, select: u.id))
+  end
+
   def get_follow_requests_query(%User{} = user) do
     from(
       a in Activity,
@@ -666,7 +678,7 @@ defmodule Pleroma.User do
     Repo.all(query)
   end
 
-  def search(query, resolve \\ false) do
+  def search(query, resolve \\ false, for_user \\ nil) do
     # strip the beginning @ off if there is a query
     query = String.trim_leading(query, "@")
 
@@ -674,16 +686,28 @@ defmodule Pleroma.User do
       User.get_or_fetch_by_nickname(query)
     end
 
+    processed_query =
+      query
+      |> String.replace(~r/\W+/, " ")
+      |> String.trim()
+      |> String.split()
+      |> Enum.map(&(&1 <> ":*"))
+      |> Enum.join(" | ")
+
     inner =
       from(
         u in User,
         select_merge: %{
-          search_distance:
+          search_rank:
             fragment(
-              "? <-> (? || coalesce(?, ''))",
+              """
-              ^query,
+              ts_rank_cd(
-              u.nickname,
+                setweight(to_tsvector('simple', regexp_replace(nickname, '\\W', ' ', 'g')), 'A') ||
-              u.name
+                setweight(to_tsvector('simple', regexp_replace(coalesce(name, ''), '\\W', ' ', 'g')), 'B'),
+                to_tsquery('simple', ?)
+              )
+              """,
+              ^processed_query
             )
         },
         where: not is_nil(u.nickname)
@@ -692,11 +716,44 @@ defmodule Pleroma.User do
     q =
       from(
         s in subquery(inner),
-        order_by: s.search_distance,
+        order_by: [desc: s.search_rank],
         limit: 20
       )
 
-    Repo.all(q)
+    results =
+      q
+      |> Repo.all()
+      |> Enum.filter(&(&1.search_rank > 0))
+
+    weighted_results =
+      if for_user do
+        friends_ids = get_friends_ids(for_user)
+        followers_ids = get_followers_ids(for_user)
+
+        Enum.map(
+          results,
+          fn u ->
+            search_rank_coef =
+              cond do
+                u.id in friends_ids ->
+                  1.2
+
+                u.id in followers_ids ->
+                  1.1
+
+                true ->
+                  1
+              end
+
+            Map.put(u, :search_rank, u.search_rank * search_rank_coef)
+          end
+        )
+        |> Enum.sort_by(&(-&1.search_rank))
+      else
+        results
+      end
+
+    weighted_results
   end
 
   def blocks_import(%User{} = blocker, blocked_identifiers) when is_list(blocked_identifiers) do

lib/pleroma/web/mastodon_api/mastodon_api_controller.ex

@@ -772,7 +772,7 @@ defmodule Pleroma.Web.MastodonAPI.MastodonAPIController do
   end
 
   def search2(%{assigns: %{user: user}} = conn, %{"q" => query} = params) do
-    accounts = User.search(query, params["resolve"] == "true")
+    accounts = User.search(query, params["resolve"] == "true", user)
 
     statuses = status_search(user, query)
 
@@ -796,7 +796,7 @@ defmodule Pleroma.Web.MastodonAPI.MastodonAPIController do
   end
 
   def search(%{assigns: %{user: user}} = conn, %{"q" => query} = params) do
-    accounts = User.search(query, params["resolve"] == "true")
+    accounts = User.search(query, params["resolve"] == "true", user)
 
     statuses = status_search(user, query)
 
@@ -817,7 +817,7 @@ defmodule Pleroma.Web.MastodonAPI.MastodonAPIController do
   end
 
   def account_search(%{assigns: %{user: user}} = conn, %{"q" => query} = params) do
-    accounts = User.search(query, params["resolve"] == "true")
+    accounts = User.search(query, params["resolve"] == "true", user)
 
     res = AccountView.render("accounts.json", users: accounts, for: user, as: :user)
 

lib/pleroma/web/twitter_api/twitter_api_controller.ex

@@ -675,7 +675,7 @@ defmodule Pleroma.Web.TwitterAPI.Controller do
   end
 
   def search_user(%{assigns: %{user: user}} = conn, %{"query" => query}) do
-    users = User.search(query, true)
+    users = User.search(query, true, user)
 
     conn
     |> put_view(UserView)

test/user_test.exs

@@ -781,8 +781,7 @@ defmodule Pleroma.UserTest do
       _user_three = insert(:user, %{name: "ebn", nickname: "lain@mastodon.social"})
       user_four = insert(:user, %{nickname: "lain@pleroma.soykaf.com"})
 
-      assert user_four ==
+      assert user_four == User.search("lain@ple") |> List.first() |> Map.put(:search_rank, nil)
-               User.search("lain@ple") |> List.first() |> Map.put(:search_distance, nil)
     end
 
     test "finds a user whose name is nil" do
@@ -792,7 +791,7 @@ defmodule Pleroma.UserTest do
       assert user_two ==
                User.search("lain@pleroma.soykaf.com")
                |> List.first()
-               |> Map.put(:search_distance, nil)
+               |> Map.put(:search_rank, nil)
     end
   end