Привет, Хабр! В прошлой статье я  писал о том, как реализовать end-to-end тестирование telegram-бота. А сегодня расскажу о том, как реализовать полнотекстовый поиск в Postgres посредством SQLAlchemy и как его проиндексировать. 

Меня зовут Михаил Выборный, я python-разработчик, backend-developer в облачном провайдере beeline cloud. Вы узнаете, как подготовить файл миграции для alembic. В конце статьи я приведу ссылку на небольшой репозиторий, где реализовал сервис-класс для удобной интеграции поиска в свой проект. 

Вступление

Тема полнотекстового поиска (Full Text Search, FTS) уже не раз поднималась внутри сообщества. Конечно, для реализации полноценного сервиса, не обойтись без Sphinx или Elasticsearch, но если поиск в вашем приложении выполняет больше формальную функцию – намного удобнее реализовать FTS напрямую в БД. Не нужно поднимать еще один сторонний сервис и поддерживать в нем консистентные данные, что может быть проблемой.

Postgres предлагает несколько решений для реализации FTS.

• Операторы LIKE / ILIKE

• Поиск строки ts_query по документу ts_vector

• Расширение pg_trgm: поиск по Триграммам (Trigram / Trigraph)

Но LIKE / ILIKE не поддерживает индексы и ранжирование, а ts_query / ts_vector требует тонкой настройки словарей и становится проблемой, когда в документе содержится несколько языков. Самым простым решением из коробки будет поиск по Триграммам.

Триграмма — это группа трех последовательных символов, взятых из строки. Мы можем измерить схожесть двух строк, подсчитав число триграмм, которые есть в обеих. Эта простая идея оказывается очень эффективной для измерения схожести слов на многих естественных языках.

Таким образом pg_trgm поддерживает Fuzzy Text Search (нечеткий поиск с опечатками) в том числе и на Русском языке.

Поиск pg_trgm в связке с SQLAlchemy

Подготовка к поиску

• Для простоты примера реализуем поиск по такой таблице.

class Base(MappedAsDataclass, DeclarativeBase):
	id: Mapped[UUID] = mapped_column(primary_key=True, init=False, default_factory=uuid.uuid4)

class Article(Base, kw_only=True):
	__tablename__ = 'articles'

	title: Mapped[str | None]
	body: Mapped[str | None]

• Установим расширение.

session.execute(text('CREATE EXTENSION IF NOT EXISTS pg_trgm'))

• Добавим тестовую запись.

session.add(
	Article(
		title='Full Text Search', 
		body='Full Text Search in PostgreSQL with SQLAlchemy. '
	)
)

• Для отладки может быть полезно посмотреть, какие триграммы содержатся в записях.

session.execute(select(func.show_trgm(Article.title))).all()

# получим следующие триграммы для поля title
['  f', '  s', '  t', ' fu', ' se', ' te', 'arc', 'ch ', 'ear', 
 'ext', 'ful', 'll ', 'rch', 'sea', 'tex', 'ull', 'xt ']

• Так же мы можем посмотреть similarity_ratio для нашего поискового запроса. Этот коэффициент указывает на сходство двух строк между собой.

session.execute(select(func.similarity(Article.title, 'search string')).all()

• По умолчанию в поисковую раздачу попадают все строки, с коэффициентом similarity больше 0.3. Установим более мягкий лимит для текущей сессии.

session.execute(select(func.set_limit('0.01')))

Поисковый запрос

Для полнотекстового поиска SQLAlchemy предлагает универсальный метод match . Для Postgres это будет оператор @@, который реализует поиск ts_query по ts_vector.

select(Article.body.match('search string')) # поиск ts_query по ts_vector

Но для поиска по Триграммам отдельного метода в Алхимии нет, поэтому будем использовать булевый оператор pg_trgm: знак процента.

Article.body.bool_op('%')('search string')

Целиком запрос будет выглядеть следующим образом.

term = 'search string'
select(
	Article.body,

	# добавим в выдачу коэффициент совпадения для отладки
	func.similarity(Article.body, term), 
)
.where(
	# все строки, для которых коэффициент совпадения больше текущего лимита
	Article.body.bool_op('%')(term),
)
.order_by(
	# добавим ранжирование по коэффициенту совпадения
	func.similarity(Article.body, term).desc(), 
)

Поиск по нескольким колонкам

Внимательный читатель заметил, что мы реализовали поиск только по одной колонке body. Для поиска по нескольким полям воспользуемся методом concat.

columns = func.coalesce(Article.title, '').concat(func.coalesce(Article.body, ''))
columns = columns.self_group() # оборачивает конкатинацию столбцов в "скобки"

???? Помимо concat, необходимо обернуть каждое поле в coalesce, чтобы в случае значения None одного из полей, мы не теряли запись целиком. Подробнее про concat и оператор || в Postgres можно прочитать тут.

Теперь мы можем реализовать поиск по нескольким полям.

term = 'search string'
select(
	Article.body,
	Article.title,
	func.similarity(columns, term),  # используем конкатенированные столбцы
)
.where(
	columns.bool_op('%')(term),      # здесь тоже
)
.order_by(
	func.similarity(columns, term).desc(),  # и здесь
)

Недостатки pg_trgm

Но не стоит забывать и о минусах данного решения. Для меня это не всегда ожидаемое ранжирование поисковой выдачи. К примеру для поиска фразы Full Text Search по следующим статьям:

Article(title='Full Text Search. ', body='Full Text Search in PostgreSQL with SQLAlchemy. '),
Article(title='Full Text Search. ', body=None)

Мы получим этот результат:

('Full Text Search. ', None, 0.4166666567325592),
('Full Text Search. ', 'Full Text Search in PostgreSQL with SQLAlchemy. ', 0.22727273404598236),

Обратите внимание, в первой записи фраза 'Full Text Search' повторяется дважды, но она получила меньший similarity_ratio. Это вызвано тем, что в первой записи поля text и body имеют большее количество символов, а значит большее количество триграмм в том числе, что понижает итоговый коэффициент совпадения.

Если вам необходим "умный" поиск, где возможно настроить поисковый запрос, установить разный вес для отдельных полей (когда совпадение в заголовке ценится больше, чем в теле статьи), регулировать удаленность слов друг от друга, то стоит посмотреть в сторону поиска по ts_query / ts_vector. Так же советую обратить внимание на индекс RUM и словарь hunspell от Postgres PRO для реализации такого поиска.

Индекс

Построение индекса

Postgres поддерживает функциональные индексы, что позволяет нам построить индекс для поиска по Триграммам. В SQLAlchemy это делается следующим образом.

Index(
	'article_title_body_trgm_idx',

	# индекс необходимо строить по вызову точно такой же функции,
	# что будет использована при поисковом запросе
	func.coalesce(Article.title, '').concat(func.coalesce(Article.body, '')).label('columns'),
	
	# так же для индекса GIN нужно указать параметр gin_trgm_ops, 
	# так Postgres поймет,что индекс нужно строить по Триграммам
	postgresql_using='gin',
	postgresql_ops={
		'columns': 'gin_trgm_ops',
	},
)

Тестирование индекса

Чтобы протестировать поиск на небольших объемах, можно заставить Postgres воспользоваться индексом. Но не стоит этого делать в продакшене, так как Postgres имеет собственные механизмы для оптимизации запросов и иногда дешевле реализовать SeqScan, даже при наличии построенного индекса.

session.execute(text('SET enable_seqscan = OFF')) # Только для локального тестирования

Теперь можем вывести Query Plan. Для этого воспользуемся классом explain из официальной Wiki Алхимии: Query Plan SQL construct.

statement = select(...)
session.execute(explain(statement)).scalars()
Sort  (cost=42.40..42.42 rows=8 width=68)
  Sort Key: (similarity(((COALESCE(title, ''::character varying))::text || (COALESCE(body, ''::character varying))::text), '21345'::text)) DESC
  ->  Bitmap Heap Scan on article  (cost=32.07..42.28 rows=8 width=68)
        Recheck Cond: (((COALESCE(title, ''::character varying))::text || (COALESCE(body, ''::character varying))::text) % '21345'::text)
        ->  Bitmap Index Scan on article_title_body_trgm_idx  (cost=0.00..32.06 rows=8 width=0)
              Index Cond: (((COALESCE(title, ''::character varying))::text || (COALESCE(body, ''::character varying))::text) % '21345'::text)

Как мы видим из плана, запрос пойдет по индексу: Bitmap Index Scan. Если при построении функционального индекса будет другая конкатенация столбцов, Postgres не сможет использовать индекс и реализует запрос через Seq Scan.

Вы можете создать индекс сами.

# при создании всех таблиц
Base.metadata.create_all(bind=engine)

# или напрямую
index = Index('article_title_body_trgm_idx', ...)
index.create(engine)

Но это может быть неудобно для миграций схемы базы данных посредством Alembic. В таком случае лучше добавить создание индекса в ревизию.

index = Index('article_title_body_trgm_idx', ...)
...

# файл миграции
def upgrade() -> None:
	op.create_table('article', ...)
	op.create_index(
		index.name,
		index.table.name,
		index.expressions,
		postgresql_using='gin',
		postgresql_ops={'columns': 'gin_trgm_ops'},
	)
  
def downgrade() -> None:
	op.drop_index(
		'article_title_body_trgm_idx',
		table_name='article',
		postgresql_using='gin',
		postgresql_ops={'columns': 'gin_trgm_ops'},
	)
	op.drop_table('article')
	...

???? При попытке автоматически создать файл миграции через alembic revision --autogenerate индекс строится некорректно, поэтому его необходимо добавить самостоятельно.

Поиск по нескольким таблицам

На текущем примере мы разобрали, как строить поисковый запрос и создавать индекс на основе pg_trgm для нескольких столбцов одной таблицы. Но, зачастую, логика приложения требует искать по нескольким связанным таблицам сразу. Предположим, наша схема выглядит так.

class Author(Base, kw_only=True):
	__tablename__ = 'authors'

	username: Mapped[str]	
	first_name: Mapped[str | None]
	last_name: Mapped[str | None]
	
	# relations:
	articles: Mapped[list[Article]] = relationship(
		back_populates='author',
		default_factory=list,
	)

class Article(Base, kw_only=True):
	__tablename__ = 'articles'

	title: Mapped[str | None]
	body: Mapped[str | None]

	# relations:	
	author_id: Mapped[uuid.UUID] = mapped_column(ForeignKey(Author.id), init=False, repr=False)
	author: Mapped[Author] = relationship(
		back_populates='articles',
		init=False,
	)

Естественно, что мы хотим искать сразу и по имени автора и по содержанию статьи. Получать в выдачу все статьи автора, если совпало его имя и отдельные статьи других авторов, в которых упоминается это же имя.

Чтобы построить индекс по нескольким таблицам, нужно создать Matherizlized View. Использовать ее для поискового запроса. Обновлять ее при изменении данных. И подготовить миграции для Alembic. Но об этом я расскажу в следующей части.

Заключение

Несмотря на не самую очевидную связку между PostgreSQL и SQLAlchemy, можно достаточно быстро получить рабочий инструмент без имплементации сторонних сервисов и настройки связанности. По сути, все что нужно сделать:

• установить расширение pg_trgm

• добавить индекс с конкатенацией необходимых столбцов

• установить similarity_ratio для текущей сессии

• сделать запрос по этим же полям с фильтром where(columns.bool_op('%')(term))

Такая реализация FTS может быть очень полезна для MVP-решений:

• Поиск быстрый, так как использует индекс.

• Поиск разрешает опечатки, так как использует Триграммы.

• Поиск работает на любых языках, так как не требует установки словарей.

Так же, для удобства, я собрал небольшой репозиторий, где описал класс FuzzySearchService. Вы можете посмотреть реализацию в деталях или интегрировать его в в своей проект напрямую. Буду рад услышать в комментариях ваш опыт настройки FTS в своем проекте и с какими проблемами вы столкнулись.

beeline cloud — secure cloud provider. Разрабатываем облачные решения, чтобы вы предоставляли клиентам лучшие сервисы