perf(backend): search stays in tens of milliseconds at archive size for rare words, channels, deleted and transcripts

This commit is contained in:
hh committed 2026-10-19 03:41:49 +00:00
1 parent 9a6bd8a714
commit e5a5edcaae
2 files changed
+48 -29

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:
+36 -27
View File
@@ -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]