Полнотекстовый поиск в PostgreSQL с SQLAlchemy

Привет, Хабр! В прошлой статье я  писал о том, как реализовать 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() # оборачивает конкатинацию столбцов в "скобки"