Пишем движок полнотекстового поиска на Go

Полнотекстовый поиск — один из тех инструментов, которые мы используем практически каждый день, когда ищем какую-то информацию в интернете. Full-Text Search (FTS) — это метод поиска текста в коллекции документов. Документ может ссылаться на веб-страницу, газетную статью, сообщение электронной почты или любой структурированный текст.

Сегодня мы собираемся написать собственный движок FTS. К концу этой статьи он сможет выполнять поиск по миллионам документов менее чем за миллисекунду. Начнём с простых поисковых запросов, таких как «Выдать все документы со словом cat», а потом расширим движок для поддержки более сложных логических запросов.

Примечание: самым известным движком полнотекстового поиска является Lucene (а также Elasticsearch и Solr, построенные на его основе).

Зачем нужен FTS

Перед тем, как писать код, вы можете спросить: «А нельзя ли просто использовать grep или цикл с проверкой каждого документа на вхождение искомого слова?» Да, можно. Но это не всегда лучшая идея.

Корпус

Будем искать фрагменты аннотаций из англоязычной Википедии. Последний дамп доступен по адресу dumps.wikimedia.org. На сегодняшний день размер файла после распаковки составляет 913 МБ. В XML-файле более 600 тыс. документов.

Пример документа:

<title>Wikipedia: Kit-Cat Klock</title>
<url>https://en.wikipedia.org/wiki/Kit-Cat_Klock</url>
<abstract>The Kit-Cat Klock is an art deco novelty wall clock shaped like a grinning cat with cartoon eyes that swivel in time with its pendulum tail.</abstract>

Загрузка документов

Сначала нужно загрузить все документы из дампа, используя очень удобный встроенный пакет encoding/xml:

import (
    "encoding/xml"
    "os"
)

type document struct {
    Title string `xml:"title"`
    URL   string `xml:"url"`
    Text  string `xml:"abstract"`
    ID    int
}

func loadDocuments(path string) ([]document, error) {
    f, err := os.Open(path)
    if err != nil {
        return nil, err
    }
    defer f.Close()

    dec := xml.NewDecoder(f)
    dump := struct {
        Documents []document `xml:"doc"`
    }{}
    if err := dec.Decode(&dump); err != nil {
        return nil, err
    }

    docs := dump.Documents
    for i := range docs {
        docs[i].ID = i
    }
    return docs, nil
}

Каждому документу присваивается уникальный ID. Для простоты первому загруженному документу присваивается ID=0, второму ID=1 и так далее.

Первая попытка

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

Теперь у нас все документы загружены в память, попробуем найти те, в которых упоминаются кошки. Сначала пройдёмся по всем документам и проверим их на подстроку cat:

func search(docs []document, term string) []document {
    var r []document
    for _, doc := range docs {
        if strings.Contains(doc.Text, term) {
            r = append(r, doc)
        }
    }
    return r
}

На моём ноутбуке поиск занимает 103 мс — не так уж плохо. Если выборочно проверить несколько документов из выдачи, то можно заметить, что функция выдаёт соответствие на слова caterpillar и category, но не на Cat с заглавной буквой C. Это не совсем то, что мы ищем.

Прежде чем продолжать, нужно исправить две вещи:

• Сделать поиск нечувствительным к регистру (чтобы выдача включала и Cat).
  
• Учесть границы слов, а не подстроки (чтобы в выдаче не было слов вроде caterpillar и communication).

Поиск с помощью регулярных выражений

Одно из очевидных решений, которое решает обе проблемы — регулярные выражения.

В данном случае нам нужно (?i)\bcat\b:

• (?i) означает нечувствительность регулярного выражения к регистру
  
• \b указывает на соответствие границам слов (место, где с одной стороны есть символ, а с другой стороны нет)

Но теперь поиск занял больше двух секунд. Как видите, система начала тормозить даже на скромном корпусе из 600 тыс. документов. Хотя такой подход легко реализовать, он не очень хорошо масштабируется. По мере увеличения набора данных нужно сканировать всё больше документов. Временнáя сложность такого алгоритма линейна, то есть количество документов для сканирования равно общему количеству документов. Если бы у нас было 6 миллионов документов вместо 600 тысяч, поиск занял бы 20 секунд. Придётся придумать что-то получше.

Инвертированный индекс

Чтобы ускорить поисковые запросы, мы предварительно обработаем текст и построим индекс.

Ядром FTS является структура данных, которая называется инвертированный индекс. Он связывает каждое слово с документами, содержащими это слово.

Пример:

documents = {
    1: "a donut on a glass plate",
    2: "only the donut",
    3: "listen to the drum machine",
}

index = {
    "a": [1],
    "donut": [1, 2],
    "on": [1],
    "glass": [1],
    "plate": [1],
    "only": [2],
    "the": [2, 3],
    "listen": [3],
    "to": [3],
    "drum": [3],
    "machine": [3],
}

Ниже приведён реальный пример инвертированного индекса. Это указатель в книге, где термин сопровождается номерами страниц:

Анализ текста

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

Анализатор текста состоит из токенизатора и нескольких фильтров.

Токенизатор

Токенизатор — это первый шаг в анализе текста. Его задача — преобразовать текст в список токенов. Наша реализация разбивает текст на границах слов и удаляет знаки препинания:

func tokenize(text string) []string {
    return strings.FieldsFunc(text, func(r rune) bool {
        // Split on any character that is not a letter or a number.
        return !unicode.IsLetter(r) && !unicode.IsNumber(r)
    })
}

> tokenize("A donut on a glass plate. Only the donuts.")

["A", "donut", "on", "a", "glass", "plate", "Only", "the", "donuts"]

Фильтры

В большинстве случаев недостаточно простого преобразования текста в список токенов. Чтобы облегчить индексацию и поиск, потребуется дополнительная нормализация.

Строчные буквы

Чтобы поиск был нечувствителен к регистру, фильтр строчных букв преобразует токены в нижний регистр. Слова cAt, Cat и caT нормализуются до формы cat. Позже при обращении к индексу мы также нормализуем в нижний регистр и поисковые запросы, так что поисковый запрос cAt найдёт слово Cat.

Удаление общеупотребительных слов

Почти любой англоязычный текст содержит общеупотребительные слова, такие как a, I, The или be. Они называются стоп-словами и присутствуют почти во всех документах, так что их следует удалить.

Нет никакого «официального» списка стоп-слов. Давайте исключим топ-10 по списку OEC. Не стесняйтесь дополнять его:

var stopwords = map[string]struct{}{ // I wish Go had built-in sets.
    "a": {}, "and": {}, "be": {}, "have": {}, "i": {},
    "in": {}, "of": {}, "that": {}, "the": {}, "to": {},
}

func stopwordFilter(tokens []string) []string {
    r := make([]string, 0, len(tokens))
    for _, token := range tokens {
        if _, ok := stopwords[token]; !ok {
            r = append(r, token)
        }
    }
    return r
}

> stopwordFilter([]string{"a", "donut", "on", "a", "glass", "plate", "only", "the", "donuts"})

["donut", "on", "glass", "plate", "only", "donuts"]

Стемминг

Из-за грамматических правил в документах встречаются разные формы слов. Стемминг сводит их к основной форме. Например, fishing, fished и fisher сводятся к основной форме fish.

Реализация стемминга — нетривиальная задача, она не рассматривается в этой статье. Возьмём один из существующих модулей:

import snowballeng "github.com/kljensen/snowball/english"

func stemmerFilter(tokens []string) []string {
    r := make([]string, len(tokens))
    for i, token := range tokens {
        r[i] = snowballeng.Stem(token, false)
    }
    return r
}

> stemmerFilter([]string{"donut", "on", "glass", "plate", "only", "donuts"})

["donut", "on", "glass", "plate", "only", "donut"]

Примечание: Стеммеры не всегда работают корректно. Например, некоторые могут сократить airline до airlin.

Сборка анализатора

func analyze(text string) []string {
    tokens := tokenize(text)
    tokens = lowercaseFilter(tokens)
    tokens = stopwordFilter(tokens)
    tokens = stemmerFilter(tokens)
    return tokens
}

Токенизатор и фильтры преобразуют предложения в список токенов:

> analyze("A donut on a glass plate. Only the donuts.")

["donut", "on", "glass", "plate", "only", "donut"]

Токены готовы к индексации.

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

Вернёмся к инвертированному индексу. Он сопоставляет каждое слово с идентификаторами документов. Для хранения карты (отображения) хорошо подходит встроенный тип данных map. Ключом будет токен (строка), а значением — список идентификаторов документов:

type index map[string][]int

В процессе построения индекса происходит анализ документов и добавление их идентификаторов в карту:

func (idx index) add(docs []document) {
    for _, doc := range docs {
        for _, token := range analyze(doc.Text) {
            ids := idx[token]
            if ids != nil && ids[len(ids)-1] == doc.ID {
                // Don't add same ID twice.
                continue
            }
            idx[token] = append(ids, doc.ID)
        }
    }
}

func main() {
    idx := make(index)
    idx.add([]document{{ID: 1, Text: "A donut on a glass plate. Only the donuts."}})
    idx.add([]document{{ID: 2, Text: "donut is a donut"}})
    fmt.Println(idx)
}

Всё работает! Каждый токен в отображении ссылается на идентификаторы документов, содержащих этот токен:

map[donut:[1 2] glass:[1] is:[2] on:[1] only:[1] plate:[1]]

Запросы

Для запросов к индексу применим тот же токенизатор и фильтры, которые использовали для индексации:

func (idx index) search(text string) [][]int {
    var r [][]int
    for _, token := range analyze(text) {
        if ids, ok := idx[token]; ok {
            r = append(r, ids)
        }
    }
    return r
}

> idx.search("Small wild cat")

[[24, 173, 303, ...], [98, 173, 765, ...], [[24, 51, 173, ...]]

И теперь, наконец, мы можем найти все документы, в которых упоминаются кошки. Поиск по 600 тыс. документов занял меньше миллисекунды (18 мкс)!

При использовании инвертированного индекса временнáя сложность поискового запроса линейна по отношению к числу поисковых токенов. В приведённом выше примере запроса, кроме анализа входного текста, выполняется всего три поиска по карте.

Логические запросы

Предыдущий запрос вернул несвязанный список документов для каждого токена. Но мы обычно ожидаем, что поиск по фразе small wild cat выдаёт список результатов, которые содержат одновременно small, wild и cat. Следующий шаг — вычислить пересечение между списками. Таким образом, мы получим список документов, соответствующих всем токенам.



К счастью, идентификаторы в нашем инвертированном индексе вставляются в порядке возрастания. Поскольку ID отсортированы, можно вычислить пересечение между списками в линейном времени. Функция intersection одновременно выполняет итерацию двух списков и собирает идентификаторы, которые присутствуют в обоих:

func intersection(a []int, b []int) []int {
    maxLen := len(a)
    if len(b) > maxLen {
        maxLen = len(b)
    }
    r := make([]int, 0, maxLen)
    var i, j int
    for i < len(a) && j < len(b) {
        if a[i] < b[j] {
            i++
        } else if a[i] > b[j] {
            j++
        } else {
            r = append(r, a[i])
            i++
            j++
        }
    }
    return r
}

Обновленный search анализирует заданный текст запроса, ищет токены и вычисляет заданное пересечение между списками ID:

func (idx index) search(text string) []int {
    var r []int
    for _, token := range analyze(text) {
        if ids, ok := idx[token]; ok {
            if r == nil {
                r = ids
            } else {
                r = intersection(r, ids)
            }
        } else {
            // Token doesn't exist.
            return nil
        }
    }
    return r
}

В дампе Википедии только два документа, которые одновременно содержат слова small, wild и cat:

> idx.search("Small wild cat")

130764  The wildcat is a species complex comprising two small wild cat species, the European wildcat (Felis silvestris) and the African wildcat (F. lybica).
131692  Catopuma is a genus containing two Asian small wild cat species, the Asian golden cat (C. temminckii) and the bay cat.

Поиск работает как положено!

Кстати, я впервые узнал о катопумах, вот одна из них:

Выводы

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

Я не упомянул о многих аспектах, которые могут значительно улучшить производительность и сделать поиск удобнее. Вот некоторые идеи для дальнейших улучшений:

• Добавить логические операторы OR и NOT.
  
• Хранить индекс на диске:
  
  • Восстановление индекса при каждом перезапуске приложения занимает некоторое время.
    
  • Большие индексы могут не поместиться в памяти.
• Поэкспериментировать с памятью и оптимизированными для CPU форматами данных для хранения наборов ID. Взглянуть на Roaring Bitmaps.
  
• Индексация нескольких полей документа.
  
• Сортировать результаты по релевантности.

Весь исходный код опубликован на GitHub.