Full-Text Search: Implementando com Postgres e Django

Algum tempo atrás vi o texto "A powerful full-text search in PostgreSQL in less than 20 lines" do Leandro Proença [1] e quis implementar algo assim pra projetos que não demandam o poder de um Apache Lucene ou de um Elastic Search.

O django já possui, em seu core, uma aplicação com métodos que são utilizados apenas com o Postgres e, para a minha surpresa, todos os conceitos de full-text search já estavam disponíveis nesse app.

Restou, nesse caso, tentar reproduzir, por assim dizer, a query do texto original utilizando o ORM do django e os métodos do full-text search.

Esse texto tem, por objetivo, trazer explicações sobre como essa implementação foi feita. Fundamentalmente, esse texto será uma versão explicada dessa thread no twitter.

Mostre-me o código

Todo o código-fonte do projeto está disponível no GitHub, nesse repositório.

Disclaimer:

O código da versão desse texto está disponível na branch texto-1.

Adicionando configurações necessárias

Dentro do settings.py do projeto, precisamos adicionar a aplicação django.contrib.postgres dentro da variável de INSTALLED_APPS para que possamos utilizar as ferramentas do django próprias para o Postgres:

# ...

INSTALLED_APPS = [
    'django.contrib.admin',
    'django.contrib.auth',
    'django.contrib.contenttypes',
    'django.contrib.sessions',
    'django.contrib.messages',
    'django.contrib.staticfiles',
    'django.contrib.postgres',
]

# ...

Criando o model

Precisamos criar um model para poder utilizar os conceitos da busca dentro dele. Para simplificar, esse caso, utilizamos um model com um único campo de texto para as buscas:

class Singer(models.Model):
    name = models.CharField("Cantor", max_length=150)

    def __str__(self):
        return self.name

    class Meta:
        verbose_name = "Cantor"
        verbose_name_plural = "Cantores"

Criando uma view

Para testar os conceitos de full-text search, podemos criar uma view. Antes, é necessário dizer que nesse texto estou usando views padrão do django com templates em HTML para não adicionar mais complexidade lidando com o Rest Framework.

Podemos criar uma view que recebe uma query string para fazer a busca:

from django.shortcuts import render
from .models import Singer

def search_singer(request):
    term = request.GET.get('q')
    if term:
        # TODO: fazer busca aqui
    else:
        singers = Singer.objects.order_by("-id").all()

    context = {
        'singers': singers,
        'term': term,
    }
    return render(request, "cantor.html", context)

O template cantor.html que estou utilizando é bem simples apenas para permitir testes de forma mais fácil:

<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta http-equiv="X-UA-Compatible" content="IE=edge">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>Buscando Cantores</title>
</head>
<body>

    <div>

        <form action="">
            <input type="search" name="q" {% if term %} value="{{ term }}" {% endif %} />
            <button type="submit">Pesquisar</button>
        </form>

    </div>

    {% if singers %}
        <main>
            {% for item in singers %}
                <div>
                    <h3>{{item.name}}</h3>
                    {% if item.rank or item.similarity %}
                        <div>
                            Rank: {{item.rank}}, Similaridade: {{item.similarity}}
                        </div>
                    {% endif %}
                </div>
            {% endfor %}
        </main>
    {% endif %}

</body>
</html>

Full-Text Search

Precisamos, primeiro, criar um SearchVector (ts_vector) e um SearchQuery (tsquery). Assim:

from django.contrib.postgres.search import SearchVector, SearchQuery

# ...

vector = SearchVector("name", config="portuguese")
query = SearchQuery(term, config="portuguese")

# ...

O vector é feito assim pra utilizar a coluna "name" do model Singer. A query é feita para processar a variável term recebida no código da view acima.

O próximo ponto é criar annotations para fazer o select de campos como o to_tsvector e o ts_rank (o método .annotate do Django ORM faz o select de outros campos e agrega eles a entidade):

from django.contrib.postgres.search import SearchVector, SearchQuery, SearchRank

# ...

vector = SearchVector("name", config="portuguese")
query = SearchQuery(term, config="portuguese")
singers = Singer.objects.annotate(
    search=vector,
    rank=SearchRank(vector, query),
).filter(
    search=query
).order_by("-rank").all()

# ...

Adicionando o código dentro da view, passamos a ter:

from django.shortcuts import render
from django.contrib.postgres.search import SearchVector, SearchQuery, SearchRank
from .models import Singer

def search_singer(request):
    term = request.GET.get('q')
    if term:
        vector = SearchVector("name", config="portuguese")
        query = SearchQuery(term, config="portuguese")
        singers = Singer.objects.annotate(
            search=vector,
            rank=SearchRank(vector, query),
        ).filter(
            search=query
        ).order_by("-rank").all()
    else:
        singers = Singer.objects.order_by("-id").all()

    context = {
        'singers': singers,
        'term': term,
    }
    return render(request, "cantor.html", context)

Utilizando um pequeno grupo de dados para teste:

Podeos testar e verificar que passamos a ter uma busca funcional:

Porém, ainda temos alguns problemas, pois, por exemplo, na busca por palavras incompletas, perdemos o ranqueamento:

Nesse ponto, entra a busca por similaridade que, combinada com o Full-Text Search nos permitirá fazer uma busca mais funcional.

Busca por Similaridade

Precisamos, primeiro, adicionar a extensão pg_trgm no banco de dados. Podemos fazer isso manualmente ou podemos criar uma migration vazia e adicionar essa extensão na migration. Vou seguir pela segunda opção. Para a primeira, basta executar o comando no banco de dados:

CREATE EXTENSION pg_trgm

Para a segunda abordagem, podemos executar o comando python manage.py makemigrations nome_do_app --empty e ele criará uma -migration vazia. A partir da migration vazia, podemos adicionar o import ao CreateExtension e adicionar dentro de operations:

from django.db import migrations
from django.contrib.postgres.operations import CreateExtension


class Migration(migrations.Migration):
    dependencies = [
        ('texto', '0003_alter_feat_music'),
    ]

    operations = [
        CreateExtension("pg_trgm")
    ]

Basta agora executar python manage.py migrate e teremos a extensão criada no banco de dados.

Agora, dentro da nossa busca, podemos fazer o uso do TrigramSimilarity para melhorar nossos resultados. Primeiro, vamos adicionar dentro do .annotate:

from django.contrib.postgres.search import SearchVector, SearchQuery, SearchRank, TrigramSimilarity

# ...

singers = Singer.objects.annotate(
    search=vector,
    rank=SearchRank(vector, query),
    similarity=TrigramSimilarity("name", term),
)

# ...

Precisamos, também, alterar o .filter para utilizar de um operador lógico OU. Para isso, precisamos fazer uso do Q(condição 1) | Q(condição 2) do django:

from django.contrib.postgres.search import SearchVector, SearchQuery, SearchRank, TrigramSimilarity
from django.db.models import Q

# ...

singers = Singer.objects.annotate(
    search=vector,
    rank=SearchRank(vector, query),
    similarity=TrigramSimilarity("name", term),
).filter(
    Q(search=query) | Q(similarity__gt=0)
).order_by("-rank", "-similarity").all()

# ...

Aqui, o que fazemos é adicionar o campo de similarity na nossa query e filtrar pra "o full-text search encontrou" ou "a similaridade é maior que zero". A partir desse momento, fazendo a mesma busca de um dos prints acima:

Por fim, nossa view passa a ter o código:

from django.shortcuts import render
from django.db.models import Q
from django.contrib.postgres.search import (
    SearchQuery,
    SearchRank,
    SearchVector,
    TrigramSimilarity,
)
from .models import Singer

def search_singer(request):
    term = request.GET.get('q')
    if term:
        vector = SearchVector("name", config="portuguese")
        query = SearchQuery(term, config="portuguese")
        singers = Singer.objects.annotate(
            search=vector,
            rank=SearchRank(vector, query),
            similarity=TrigramSimilarity("name", term),
        ).filter(
            Q(search=query) | Q(similarity__gt=0)
        ).order_by("-rank", "-similarity").all()
    else:
        singers = Singer.objects.order_by("-id").all()

    context = {
        'singers': singers,
        'term': term,
    }
    return render(request, "cantor.html", context)

É possível utilizar tanto o rank ou o similarity para cortar valores, conforme exemplos da documentação.

Por último, podemos adicionar um índice dentro do nosso model para lidar com performance das queries:

from django.db import models
from django.contrib.postgres.indexes import GinIndex
from django.contrib.postgres.search import SearchVector

class Singer(models.Model):
    name = models.CharField("Cantor", max_length=150)

    def __str__(self):
        return self.name

    class Meta:
        verbose_name = "Cantor"
        verbose_name_plural = "Cantores"
        indexes = [
            GinIndex(
                SearchVector("name", config="portuguese"),
                name="singer_search_vector_idx",
            )
        ]

Todo o código-fonte do projeto está disponível no GitHub, nesse repositório.

Disclaimer:

O código da versão desse texto está disponível na branch texto-1.

Referências

1 - A powerful full-text search in PostgreSQL in less than 20 lines

2 - Full text search - Django Documentation

Foto de capa por Mick Haupt no Unsplash.

Blog