From dc45ec62c2f5dfcc895854dfbddf6fe9621d3072 Mon Sep 17 00:00:00 2001 From: Ivan Tashkinov Date: Mon, 14 Jan 2019 20:04:45 +0300 Subject: [PATCH] [#477] User search improvements: tsquery search with field weights, friends & followers boosting. --- lib/pleroma/user.ex | 75 ++++++++++++++++--- .../mastodon_api/mastodon_api_controller.ex | 6 +- .../web/twitter_api/twitter_api_controller.ex | 2 +- test/user_test.exs | 5 +- 4 files changed, 72 insertions(+), 16 deletions(-) diff --git a/lib/pleroma/user.ex b/lib/pleroma/user.ex index 681280539..52638b446 100644 --- a/lib/pleroma/user.ex +++ b/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, - u.nickname, - u.name + """ + ts_rank_cd( + setweight(to_tsvector('simple', regexp_replace(nickname, '\\W', ' ', 'g')), 'A') || + 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 diff --git a/lib/pleroma/web/mastodon_api/mastodon_api_controller.ex b/lib/pleroma/web/mastodon_api/mastodon_api_controller.ex index a8fe9d708..54367f586 100644 --- a/lib/pleroma/web/mastodon_api/mastodon_api_controller.ex +++ b/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) diff --git a/lib/pleroma/web/twitter_api/twitter_api_controller.ex b/lib/pleroma/web/twitter_api/twitter_api_controller.ex index 1c728166c..ede079963 100644 --- a/lib/pleroma/web/twitter_api/twitter_api_controller.ex +++ b/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) diff --git a/test/user_test.exs b/test/user_test.exs index cfccce8d1..efa7937bc 100644 --- a/test/user_test.exs +++ b/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 == - User.search("lain@ple") |> List.first() |> Map.put(:search_distance, nil) + assert user_four == User.search("lain@ple") |> List.first() |> Map.put(:search_rank, 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