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

BEIR — это бенчмарки поиска информации, ориентированные на сценарии использования в формате «вопрос-ответ».

Мой хобби-проект SearchArray добавляет в Pandas полнотекстовый поиск. Поэтому естественно, чтобы ощутить трепет от моих потрясающих навыков разработчика, я решил использовать BEIR для сравнения SearchArray с Elasticsearch (с тем же запросом + токенизацией). Поэтому я потратил субботу на интеграцию SearchArray в BEIR и измерение релевантности и производительности с корпусом MSMarco Passage Retrieval (8 миллионов документов).

Барабанная дробь...

… печальные звуки тромбона.

Моя библиотека хуже во всех отношениях.

По крайней мере, NDCG@10 почти равны, а значит, я правильно реализовал вычисление BM25 (вероятно, из-за незначительной разницы в токенизации)

Что по синдрому самозванца?

Я не утопаю в своём стыде, а ТОЧНО знаю, что случилось… И история эта весьма поучительна. Давайте поговорим о том, почему быстр реальный поисковый движок, а не хобби-проект.

Волшебный WAND

(Или почему SearchArray выполняет поиск верхних 8 миллионов документов, а Elasticsearch — поиск верхних K)

В системах лексического поиска выполняется поиск множественных терминов. Мы берём оценку BM25 каждого термина, а затем комбинируем все эти оценки в окончательную оценку документа. Например, если поиск luke skywalker на самом деле означает BM25(luke) ??? BM25(skywalker), где ??? — некий математический оператор.

В простом запросе OR мы просто берём сумму (SUM) каждого термина для каждого документа, например, в поиске luke skywalker оценка BM25(luke) + BM25(skywalker) может выглядеть так:

SearchArray выполняет лишь вычисление оценки BM25. От него можно получать большие массивы numpy оценок BM25 каждого документа. Затем мы комбинируем оценки, буквально воспользовавшись np.sum. Разумеется, поисковый движок наподобие Elasticsearch ведёт себя иначе. У него есть другая гарантия — для запроса OR он получает N документов с наивысшей оценкой.

Это небольшое отличие обеспечивает поисковым движкам большую долю свободы. Поисковые движки могут использовать алгоритм Weak-AND (WAND) , чтобы избежать дополнительной работы при комбинировании оценок множественных терминов в окончательные результаты лучших N.

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

Системы оценки наподобие BM25 сильно зависят от частоты термина в документах. Поэтому редкие термины с высоким (1 / частотность в документах) имеют большую вероятность повлиять на конечную оценку и оказаться в верхних K. К счастью, эти термины (например, skywalker) встречаются в меньшем количестве документов. Поэтому мы можем получать эти избранные, элитарные немногочисленные документы в структуре данных, выполняющей сопоставление skywalker -> [... несколько id соответствующих документов...] (также называемой posting). Мы можем глубоко исследовать этот список.

С другой стороны, мы можем гораздо осмотрительнее обращаться со скучными распространёнными терминами наподобие luke. И это полезно, потому что список postings термина luke очень затратен luke -> [... огромный список документов ...]. Мы бы предпочли не сканировать всё это.

Мы можем представить, что эти списки id документов наряду с их частотностью терминов будут ещё одним важным источником входных данных BM25. А если они отсортированы в порядке убывания частотности терминов, то то можно спускаться вниз по этому списку, пока оценка BM25 термина больше не сможет попасть в верхние K результатов, после чего выполнить ранний выход.

WAND и подобные ему оптимизации помогают Elasticsearch избегать лишней работы, однако SearchArray с радостью выполняет её, возвращая огромный дурацкий массив оценок BM25.

Если посмотреть на этот icicle-график производительности SearchArray при выполнении поиска OR, можно увидеть, что всё время тратится на суммирование огромного массива и необязательное вычисление оценок BM25 для множества документов.

SearchArray не сохраняет postings напрямую

В отличие от большинства поисковых движков SearchArray не хранит списки postings «термины -> документы».

Вместо этого SearchArray хранит позиционный индекс, создаваемый в первую очередь для сопоставления фраз. Мы передаём SearchArray список терминов ['mary', 'had', 'a', 'little', 'lamb']. Затем он находит все места, где mary встречается на одну позицию раньше had и так далее. Он реализует это, сохраняя для каждого термина позиции в виде roaring bitmap.

В нашей roaring bitmap каждое 64-битное слово имеет заголовок, указывающий на расположение позиций (документ и область документа). Каждая позиция бита соответствует позиции в документе. 1 обозначает, что термин присутствует, 0 — что он отсутствует.

Поэтому чтобы создать коллекцию сопоставления фраз для mary had, мы можем просто находить места, в которых биты одного термина находятся рядом с битами другого. Это очень быстро можно выполнять при помощи битовой арифметики.

mary
 <Roaring Headers (doc id, etc)>  00000010000000    | 00000000000010
had
 <Roaring Headers (doc id, etc)>  00000001000000    | 00000000000001      #< Сдвиг влево на 1, AND, считаем биты для получения частотности фразы
...

Но удобное свойство такого подхода, к тому же упрощающее поддержку этого хобби-проекта, заключается в том, что можно использовать его и для вычисления частотности терминов. Просто выполнив popcount (подсчёт количества ненулевых битов), а затем получив эти документы для термина, мы получаем отображение «id документов -> частотность термина».

Поэтому, как видно ниже, мы тратим на это приличный объём времени:

На самом деле я соврал: хоть это и базовый механизм для хранения частотности терминов, мы применяем кэширование. Кэш запоминает отображение «id документа -> частотности терминов», когда массив roaring bit > N 64-битных слов. Это позволяет пользователям настраивать N для балансирования памяти и скорости, а также приблизиться по скорости к списку postings.

Кэширование независящих от терминов компонентов BM25

Посмотрите на этот фрагмент кода для вычисления BM25:

        bm25 = (
            term_freq / (term_freq + (k1 * ((1 - b) + (b * (doc_len / avg_doc_lens)))))
        ) * idf

Здесь стоит обратить внимание на то, что БОЛЬШАЯ ЧАСТЬ этого вычисления никак не связана с терминами, которые мы ищем:

 (k1 * ((1 - b) + (b * (doc_len / avg_doc_lens)))))

В моих тестах выяснилось, что часть задержек (1 мс для 1 миллиона документов) можно снизить, кэшировав куда-нибудь всё в k1 + ... avg_doc_lens. Если doc_len соответствует массиву со значением длины документов для каждого документа, можно создать массив с кэшированием этой формулы. Но немного сложнее будет поддерживать этот один дополнительный глобальный кэш, поэтому я пока его не использовал.

Кэширование полного запроса, а не только отдельных оценок терминов BM25

SearchArray — это всего лишь система для вычисления оценок BM25 (или любой другой степени похожести). Можно использовать её для создания «запроса or» или чего-то подобного при помощи… Сама она этого не делает. Например, реализованный в BEIR код выглядит так:

def bm25_search(corpus, query, column):
    tokenizer = corpus[column].array.tokenizer
    query_terms = tokenizer(query)
    scores = np.zeros(len(corpus))
    query_terms = set(query_terms)
    for term in query_terms:
            scores += corpus[column].array.score(term)
    return scores

Но в обычном поисковом движке наподобие Solr, Elasticsearch, OpenSearch или Vespa эта логика выражается в Query DSL движка. Поэтому поисковый движок может выполнять планирование+кэширование полных вычислений, а SearchArray вручает вам все инструменты, чтобы, с точки зрения производительности, выстрелить себе в ногу (не говоря уже о том, что говорилось выше про WAND и тому подобном).

Именно поэтому мы должны быть благодарны людям, разрабатывающим поиск

SearchArray — это инструмент для прототипирования при помощи обычного инструментария Pydata, а не для создания огромных поисковых систем наподобие Elasticsearch. Полезно понимать компромиссы, используемые при создании лексической системы, потому что эти системы делают упор на разное.

Было бы здорово, если бы мы могли выражать наши запросы в таких ориентированных на кадры данных DSL. Например, в ленивой системе поиска верхних N наподобие Polars, получающей данные из различных источников, вычисляющей их оценку и суммирование, после чего выполняющей любые произвольные вычисления с оценками. Я могу только надеяться, что когда-нибудь что-то подобное появится. Пока разработчики создают эти DAG менее выразительным образом: в рамках своей модели Torch DAG или в какой-нибудь самодельной системе DAG времени запросов.

Как бы то ни было, я потрясён людьми, которые занимаются большими крупномасштабными распределёнными лексическими поисковыми системами наподобие Vespa, Lucene, OpenSearch, Elasticsearch, Solr. Эти люди должны стать героями и для вас: они выполняют за всех нас эту трудную работу, и мы НЕ должны воспринимать её, как что-то само собой разумеющееся.

Ниже представлены примечания и приложения по BEIR и различным скриптам бенчмаркинга.

Ссылки на скрипты

• Скрипт Elasticsearch | Конфигурация Elasticsearch

• Скрипт SearchArray

Приложение: интеграция с BEIR…

BEIR предоставляет набор встроенных датасетов и инструментов получения метрик, если реализовать класс BaseSearch со следующей сигнатурой:

    class SearchArraySearch(BaseSearch):

        def search(self,
                   corpus: Dict[str, Dict[str, str]],
                   queries: Dict[str, str],
                   top_k: int,
                   *args,
                   **kwargs) -> Dict[str, Dict[str, float]]:

Входные данные:

• Corpus: словарь, указывающий id документа на множество полей для индексации, например

{'text': 'The presence of communication amid scientific minds was equally important to the success of the Manhattan Project as scientific intellect was. The only cloud hanging over the impressive achievement of the atomic researchers and engineers is what their success truly meant; hundreds of thousands of innocent lives obliterated.',
 'title': ''}
...

• Queries: словарь, указывающий «id запроса -> запрос»:

{"1": "Who was the original governor of the plymouth colony"}
...

Вывод — это словарь запросов «ids -> {doc ids -> scores}», где каждый запрос имеет оценку top_k.

При вызове поиска необходимо:

1. Индексировать корпус

2. Передать все запросы и собрать оценки

По сути, это выглядит примерно так:

def search(self,
           corpus: Dict[str, Dict[str, str]],
           queries: Dict[str, str],
           top_k: int,
           *args,
           **kwargs) -> Dict[str, Dict[str, float]]:
    corpus = self.index_corpus(corpus)     # <- добавляем токенизированные столбцы к кадру данных
    results = {}

    for query_id, query in queries.items():    #<- ищем и получаем результаты
        results_for_query = some_search_function(corpus, query)
        results[query_id] = get_top_k(results_for_query)
    return results

Как это выглядит для SearchArray?

Для индексации мы обходим в цикле каждый столбец str и добавляем столбец SearchArray в DF. Ниже показана токенизация при помощи токенизайтора snowball:

            for column in corpus.columns:
                if corpus[column].dtype == 'object':
                    corpus[column].fillna("", inplace=True)
                    corpus[f'{column}_snowball'] = SearchArray.index(corpus[column],
                                                                     data_dir=DATA_DIR,
                                                                     tokenizer=snowball_tokenizer)

Затем мы заменяем показанную выше функцию some_search_function на что-то, что выполняет поиск по столбцам SearchArray. Например, простым bm25_search:

def bm25_search(corpus, query):
    query = snowball_tokenizer(query)
    scores = np.zeros(len(corpus))
    for q in query:
        scores += corpus['text_snowball'].array.score(q)
    return scores

(Я не стал описывать возню с потоками. Полный код можно посмотреть на Github.)