102 lines
3.6 KiB
Go
102 lines
3.6 KiB
Go
package search
|
|
|
|
import (
|
|
"context"
|
|
"math"
|
|
"sort"
|
|
|
|
"github.com/google/uuid"
|
|
"github.com/jackc/pgx/v5/pgxpool"
|
|
)
|
|
|
|
// suggestPool — во сколько раз больше кандидатов берётся из Meilisearch для переранжирования подсказок.
|
|
const suggestPool = 3
|
|
|
|
// Exact names stay ahead of related artists, regardless of listening history.
|
|
func personalizeArtists(found, exact []string, aff map[uuid.UUID]float64, n int) []string {
|
|
protected := make(map[string]bool, len(exact))
|
|
for _, id := range exact {
|
|
protected[id] = true
|
|
}
|
|
var first, rest []string
|
|
for _, id := range found {
|
|
if protected[id] {
|
|
first = append(first, id)
|
|
} else {
|
|
rest = append(rest, id)
|
|
}
|
|
}
|
|
first = append(first, personalize(rest, aff, max(0, n-len(first)))...)
|
|
return first[:min(n, len(first))]
|
|
}
|
|
|
|
// Affinity — насколько пользователь «знает» исполнителей, альбомы и треки: прослушивания, доигрывания,
|
|
// лайки и дизлайки. Отрицательное — не любит. Нужен для персональных подсказок поиска.
|
|
type Affinity struct{ pool *pgxpool.Pool }
|
|
|
|
// NewAffinity создаёт источник предпочтений.
|
|
func NewAffinity(pool *pgxpool.Pool) *Affinity { return &Affinity{pool: pool} }
|
|
|
|
// Of возвращает вес для каждого из ids, о котором что-то известно.
|
|
func (a *Affinity) Of(ctx context.Context, user uuid.UUID, ids []uuid.UUID) (map[uuid.UUID]float64, error) {
|
|
rows, err := a.pool.Query(ctx, `
|
|
WITH s AS (
|
|
SELECT track_id, sum(w) AS w FROM (
|
|
SELECT track_id, (play_count + 2 * complete_count - early_skip_count)::float8 AS w
|
|
FROM user_track_stats WHERE user_id = $1
|
|
UNION ALL
|
|
SELECT track_id, CASE reaction WHEN 'like' THEN 10 ELSE -10 END FROM track_reactions WHERE user_id = $1
|
|
) x GROUP BY track_id
|
|
)
|
|
SELECT track_id, w FROM s WHERE track_id = ANY($2)
|
|
UNION ALL
|
|
SELECT t.album_id, sum(s.w) FROM s JOIN tracks t ON t.id = s.track_id WHERE t.album_id = ANY($2) GROUP BY t.album_id
|
|
UNION ALL
|
|
SELECT ta.artist_id, sum(s.w) FROM s JOIN track_artists ta ON ta.track_id = s.track_id
|
|
WHERE ta.artist_id = ANY($2) GROUP BY ta.artist_id
|
|
UNION ALL
|
|
SELECT artist_id, CASE reaction WHEN 'like' THEN 20 ELSE -20 END FROM artist_reactions
|
|
WHERE user_id = $1 AND artist_id = ANY($2)`, user, ids)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
out := map[uuid.UUID]float64{}
|
|
for rows.Next() {
|
|
var id uuid.UUID
|
|
var w float64
|
|
if err := rows.Scan(&id, &w); err != nil {
|
|
return nil, err
|
|
}
|
|
out[id] += w
|
|
}
|
|
return out, rows.Err()
|
|
}
|
|
|
|
// personalize переставляет найденное: знакомое и любимое поднимается на несколько позиций, нелюбимое
|
|
// опускается; релевантность Meilisearch остаётся основой порядка. Возвращает не больше n.
|
|
func personalize(found []string, aff map[uuid.UUID]float64, n int) []string {
|
|
type item struct {
|
|
id string
|
|
key float64
|
|
}
|
|
items := make([]item, len(found))
|
|
for i, s := range found {
|
|
key := float64(i)
|
|
if id, err := uuid.Parse(s); err == nil {
|
|
if w := aff[id]; w > 0 {
|
|
key -= 2 * math.Log1p(w)
|
|
} else if w < 0 {
|
|
key += 3
|
|
}
|
|
}
|
|
items[i] = item{s, key}
|
|
}
|
|
sort.SliceStable(items, func(i, j int) bool { return items[i].key < items[j].key })
|
|
out := make([]string, 0, min(n, len(items)))
|
|
for _, it := range items[:min(n, len(items))] {
|
|
out = append(out, it.id)
|
|
}
|
|
return out
|
|
}
|