perf(backend): search stays in tens of milliseconds at archive size for rare words, channels, deleted and transcripts
This commit is contained in:
1 parent
9a6bd8a714
commit
e5a5edcaae
2 files changed
+50
-31
No files matched your search
@@ -1,7 +1,7 @@
|
|||||||
"""message search index outside the hypertable, versions searchable
|
"""message search index outside the hypertable, versions searchable
|
||||||
|
|
||||||
Revision ID: d8f1a6c3e5b7
|
Revision ID: d8f1a6c3e5b7
|
||||||
Revises: c9e2a7f4b1d8
|
Revises: c5e7a9b1d3f6
|
||||||
Create Date: 2026-10-19 03:00:00.000000
|
Create Date: 2026-10-19 03:00:00.000000
|
||||||
|
|
||||||
"""
|
"""
|
||||||
@@ -11,7 +11,7 @@ from collections.abc import Sequence
|
|||||||
from alembic import op
|
from alembic import op
|
||||||
|
|
||||||
revision: str = "d8f1a6c3e5b7"
|
revision: str = "d8f1a6c3e5b7"
|
||||||
down_revision: str | None = "c9e2a7f4b1d8"
|
down_revision: str | None = "c5e7a9b1d3f6"
|
||||||
branch_labels: str | Sequence[str] | None = None
|
branch_labels: str | Sequence[str] | None = None
|
||||||
depends_on: str | Sequence[str] | None = None
|
depends_on: str | Sequence[str] | None = None
|
||||||
|
|
||||||
@@ -95,6 +95,13 @@ def upgrade() -> None:
|
|||||||
"CREATE INDEX ix_message_search_sender ON message_search "
|
"CREATE INDEX ix_message_search_sender ON message_search "
|
||||||
"(account_id, sender_id, date DESC)"
|
"(account_id, sender_id, date DESC)"
|
||||||
)
|
)
|
||||||
|
for flag in ("deleted", "edited"):
|
||||||
|
op.execute(
|
||||||
|
f"CREATE INDEX ix_message_search_{flag} ON message_search "
|
||||||
|
"(account_id, date DESC, chat_id DESC, message_id DESC) "
|
||||||
|
f"WHERE {flag}_at IS NOT NULL"
|
||||||
|
)
|
||||||
|
op.execute("ALTER TABLE message_search ALTER COLUMN tsv SET STATISTICS 1000")
|
||||||
op.execute(_SYNC)
|
op.execute(_SYNC)
|
||||||
op.execute(
|
op.execute(
|
||||||
"CREATE TRIGGER messages_search_insert AFTER INSERT ON messages "
|
"CREATE TRIGGER messages_search_insert AFTER INSERT ON messages "
|
||||||
@@ -117,6 +124,9 @@ def upgrade() -> None:
|
|||||||
"CREATE INDEX ix_media_tsv ON media USING gin (search_vector(extracted_text))"
|
"CREATE INDEX ix_media_tsv ON media USING gin (search_vector(extracted_text))"
|
||||||
)
|
)
|
||||||
op.execute("DROP INDEX IF EXISTS ix_messages_tsv")
|
op.execute("DROP INDEX IF EXISTS ix_messages_tsv")
|
||||||
|
for index in ("ix_message_versions_tsv", "ix_media_tsv"):
|
||||||
|
op.execute(f"ALTER INDEX {index} ALTER COLUMN 1 SET STATISTICS 1000")
|
||||||
|
op.execute("ANALYZE message_search, message_versions, media")
|
||||||
|
|
||||||
|
|
||||||
def downgrade() -> None:
|
def downgrade() -> None:
|
||||||
|
|||||||
@@ -2,6 +2,7 @@ import asyncpg
|
|||||||
|
|
||||||
from utils.search.models import (
|
from utils.search.models import (
|
||||||
MAX_LIMIT,
|
MAX_LIMIT,
|
||||||
|
ChatKind,
|
||||||
SearchFilters,
|
SearchFilters,
|
||||||
SearchHit,
|
SearchHit,
|
||||||
SearchPage,
|
SearchPage,
|
||||||
@@ -20,14 +21,14 @@ FROM translate($1, 'Ёё', 'Ее') AS q
|
|||||||
"""
|
"""
|
||||||
|
|
||||||
_CHANNELS = """
|
_CHANNELS = """
|
||||||
(SELECT chat_id FROM (
|
SELECT coalesce(array_agg(chat_id), '{}'::bigint[]) FROM (
|
||||||
SELECT DISTINCT ON (chat_id) chat_id,
|
SELECT DISTINCT ON (chat_id) chat_id,
|
||||||
coalesce(raw->'chat'->>'type', raw->>'type') AS type
|
coalesce(raw->'chat'->>'type', raw->>'type') AS type
|
||||||
FROM chat_history
|
FROM chat_history
|
||||||
WHERE account_id = {account}
|
WHERE account_id = $1
|
||||||
AND coalesce(raw->'chat'->>'type', raw->>'type') IS NOT NULL
|
AND coalesce(raw->'chat'->>'type', raw->>'type') IS NOT NULL
|
||||||
ORDER BY chat_id, ts DESC
|
ORDER BY chat_id, ts DESC
|
||||||
) typed WHERE type = 'ChatType.CHANNEL')
|
) typed WHERE type = 'ChatType.CHANNEL'
|
||||||
"""
|
"""
|
||||||
|
|
||||||
_KEY_ORDER = "s.date DESC, s.chat_id DESC, s.message_id DESC"
|
_KEY_ORDER = "s.date DESC, s.chat_id DESC, s.message_id DESC"
|
||||||
@@ -43,25 +44,27 @@ _TEXT = (
|
|||||||
"FROM message_search s WHERE {where}"
|
"FROM message_search s WHERE {where}"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
_PROBE = (
|
||||||
|
"CROSS JOIN LATERAL (SELECT * FROM message_search s "
|
||||||
|
"WHERE s.chat_id = {source}.chat_id AND s.message_id = {source}.message_id "
|
||||||
|
"AND {where} LIMIT 1) s "
|
||||||
|
)
|
||||||
|
|
||||||
_STT = (
|
_STT = (
|
||||||
f"SELECT {_COLS}, 1 AS rank, md.extracted_text, NULL::text, " # noqa: S608
|
f"SELECT {_COLS}, 1 AS rank, md.extracted_text, NULL::text, " # noqa: S608
|
||||||
"NULL::timestamptz "
|
"NULL::timestamptz FROM media md "
|
||||||
"FROM media md JOIN message_search s ON s.account_id = md.account_id "
|
f"{_PROBE.replace('{source}', 'md')}"
|
||||||
"AND s.chat_id = md.chat_id AND s.message_id = md.message_id "
|
|
||||||
"WHERE md.account_id = {account} "
|
"WHERE md.account_id = {account} "
|
||||||
"AND search_vector(md.extracted_text) @@ {tsquery} "
|
"AND search_vector(md.extracted_text) @@ {tsquery}"
|
||||||
"AND {where}"
|
|
||||||
)
|
)
|
||||||
|
|
||||||
_VERSION = (
|
_VERSION = (
|
||||||
"SELECT DISTINCT ON (s.date, s.chat_id, s.message_id) " # noqa: S608
|
"SELECT DISTINCT ON (s.date, s.chat_id, s.message_id) " # noqa: S608
|
||||||
f"{_COLS}, 2 AS rank, "
|
f"{_COLS}, 2 AS rank, NULL::text, v.text, v.observed_at "
|
||||||
"NULL::text, v.text, v.observed_at "
|
"FROM message_versions v "
|
||||||
"FROM message_versions v JOIN message_search s ON s.account_id = v.account_id "
|
f"{_PROBE.replace('{source}', 'v')}"
|
||||||
"AND s.chat_id = v.chat_id AND s.message_id = v.message_id "
|
|
||||||
"WHERE v.account_id = {account} "
|
"WHERE v.account_id = {account} "
|
||||||
"AND search_vector(v.text) @@ {tsquery} "
|
"AND search_vector(v.text) @@ {tsquery}"
|
||||||
"AND {where}"
|
|
||||||
)
|
)
|
||||||
|
|
||||||
_WITH_MEDIA = (
|
_WITH_MEDIA = (
|
||||||
@@ -83,18 +86,16 @@ class _Binder:
|
|||||||
return f"${len(self.params)}"
|
return f"${len(self.params)}"
|
||||||
|
|
||||||
|
|
||||||
def _kind(filters: SearchFilters, bind: _Binder) -> str | None:
|
def _kind(kind: ChatKind, channels: list[int], bind: _Binder) -> str:
|
||||||
if filters.kind is None:
|
if kind == "private":
|
||||||
return None
|
|
||||||
if filters.kind == "private":
|
|
||||||
return "s.chat_id > 0"
|
return "s.chat_id > 0"
|
||||||
channels = _CHANNELS.format(account=bind(filters.account_id))
|
member = f"s.chat_id = ANY({bind(channels)}::bigint[])"
|
||||||
if filters.kind == "channel":
|
return member if kind == "channel" else f"s.chat_id < 0 AND NOT {member}"
|
||||||
return f"s.chat_id IN {channels}"
|
|
||||||
return f"s.chat_id < 0 AND s.chat_id NOT IN {channels}"
|
|
||||||
|
|
||||||
|
|
||||||
def _conditions(filters: SearchFilters, bind: _Binder) -> list[str]:
|
def _conditions(
|
||||||
|
filters: SearchFilters, channels: list[int], bind: _Binder
|
||||||
|
) -> list[str]:
|
||||||
conds = [f"s.account_id = {bind(filters.account_id)}"]
|
conds = [f"s.account_id = {bind(filters.account_id)}"]
|
||||||
compared = (
|
compared = (
|
||||||
("s.chat_id =", filters.chat_id),
|
("s.chat_id =", filters.chat_id),
|
||||||
@@ -111,9 +112,8 @@ def _conditions(filters: SearchFilters, bind: _Binder) -> list[str]:
|
|||||||
for column, flag in flags
|
for column, flag in flags
|
||||||
if flag is not None
|
if flag is not None
|
||||||
]
|
]
|
||||||
kind = _kind(filters, bind)
|
if filters.kind is not None:
|
||||||
if kind is not None:
|
conds.append(_kind(filters.kind, channels, bind))
|
||||||
conds.append(kind)
|
|
||||||
if filters.cursor is not None:
|
if filters.cursor is not None:
|
||||||
date, chat_id, message_id = decode_cursor(filters.cursor)
|
date, chat_id, message_id = decode_cursor(filters.cursor)
|
||||||
conds.append(
|
conds.append(
|
||||||
@@ -128,18 +128,27 @@ def _branch(sql: str, limit: str, order: str = _KEY_ORDER, **parts: str) -> str:
|
|||||||
|
|
||||||
|
|
||||||
async def search_messages(pool: asyncpg.Pool, filters: SearchFilters) -> SearchPage:
|
async def search_messages(pool: asyncpg.Pool, filters: SearchFilters) -> SearchPage:
|
||||||
|
async with pool.acquire() as conn, conn.transaction():
|
||||||
|
await conn.execute("SET LOCAL jit = off")
|
||||||
|
return await _search(conn, filters)
|
||||||
|
|
||||||
|
|
||||||
|
async def _search(conn: asyncpg.Connection, filters: SearchFilters) -> SearchPage:
|
||||||
query = (filters.query or "").strip()
|
query = (filters.query or "").strip()
|
||||||
terms: list[str] = []
|
terms: list[str] = []
|
||||||
tsquery = None
|
tsquery = None
|
||||||
if query:
|
if query:
|
||||||
parsed = await pool.fetchrow(_PARSE, query)
|
parsed = await conn.fetchrow(_PARSE, query)
|
||||||
terms = list(parsed["terms"])
|
terms = list(parsed["terms"])
|
||||||
if not parsed["tsquery"]:
|
if not parsed["tsquery"]:
|
||||||
return SearchPage(hits=[], next_cursor=None, terms=terms)
|
return SearchPage(hits=[], next_cursor=None, terms=terms)
|
||||||
tsquery = parsed["tsquery"]
|
tsquery = parsed["tsquery"]
|
||||||
|
|
||||||
|
channels: list[int] = []
|
||||||
|
if filters.kind in {"group", "channel"}:
|
||||||
|
channels = await conn.fetchval(_CHANNELS, filters.account_id)
|
||||||
bind = _Binder()
|
bind = _Binder()
|
||||||
where = " AND ".join(_conditions(filters, bind))
|
where = " AND ".join(_conditions(filters, channels, bind))
|
||||||
size = min(filters.limit, MAX_LIMIT)
|
size = min(filters.limit, MAX_LIMIT)
|
||||||
limit = bind(size + 1)
|
limit = bind(size + 1)
|
||||||
|
|
||||||
@@ -166,7 +175,7 @@ async def search_messages(pool: asyncpg.Pool, filters: SearchFilters) -> SearchP
|
|||||||
f"ORDER BY {_KEY_ORDER}, rank LIMIT {limit}"
|
f"ORDER BY {_KEY_ORDER}, rank LIMIT {limit}"
|
||||||
)
|
)
|
||||||
|
|
||||||
rows = await pool.fetch(_WITH_MEDIA.format(hits=sql, account=account), *bind.params)
|
rows = await conn.fetch(_WITH_MEDIA.format(hits=sql, account=account), *bind.params)
|
||||||
hits = [
|
hits = [
|
||||||
SearchHit(**{**dict(row), "source": _SOURCES[row["rank"]]})
|
SearchHit(**{**dict(row), "source": _SOURCES[row["rank"]]})
|
||||||
for row in rows[:size]
|
for row in rows[:size]
|
||||||
|
|||||||
Reference in new issue
Block a user