Пользователь
0,0
рейтинг
17 июля 2012 в 15:24

Разработка → Префиксные деревья в Python

Доделал на днях питонью библиотеку datrie, реализующую префиксное дерево (см. википедию или хабр), спешу поделиться.

Если вкратце, то можно считать, что datrie.Trie — это замена стандартному питоньему dict, которая при определенных условиях (ключи — строки) занимает меньше памяти, имеет сравнимую скорость получения отдельного элемента и поддерживает дополнительные операции (получение всех префиксов данной строки, получение всех строк, начинающихся с данной строки и др.), которые работают примерно так же быстро, как и «словарные» операции.

Работает под Python 2.6-3.3, поддерживает юникод, лицензия LGPL.



datrie — это Cython-обертка над библиотекой libdatrie. В libdatrie реализован вариант префиксного дерева, в котором конечный автомат переходов между узлами хранится в специальных 2 массивах + реализовано сжатие суффиксов (ветки, которые не ветвятся, хранятся еще в одном отдельном массиве «хвостов»). Это не совсем «state of art» вариант trie (HAT-trie и тд вроде должны быть быстрее), но вариант достаточно быстрый/эффективный и с хорошей готовой реализацией (а реализация может на практике убить любой алгоритм).

Существующие варианты trie-образных структур для питона меня не устраивали. Чисто питоньи реализации неизбежно будут потреблять очень много памяти, они отметаются сразу. Другие реализации trie-образных структур для питона:

  • trie.c в biopython. Не поддерживает юникод (хотя это ОК), не работает под 3.х (и чтоб заставить работать под 3.х, нужно хорошенько переписывать обертку, т.к. библиотека реализована как «голое» C-расширение);
  • github.com/buriy/python-chartrie от buriy — довольно быстрая (на __getitem__ быстрее, чем datrie в итоге), занимает много памяти (чтоб это поправить, нужно, видимо, структуру данных менять, т.е. основу библиотеки), мало функционала, не поддерживает юникод (хотя это ОК);
  • www.dalkescientific.com/Python/PyJudy.html — старая, сложная, непонятно, работает ли (написано, что поддерживает 2.3 и 2.4)


Может еще что-то и есть, но, в любом случае, варианта «поставил — заработало — всем устраивает!» не нашел, поэтому считаю, что желание вспомнить C и поразбираться получше в Cython и Python C-API вполне успешно удалось прикрыть отсутствием нормальной реализации trie для питона.

Установка (как это обычно бывает при установке расширений, потребуется компилятор):

pip install datrie

Создание:

import string
import datrie
trie = datrie.Trie(string.ascii_lowercase)


При создании нужно сразу сказать, какие ключи можно будет с этим trie использовать (явно указывать алфавит или диапазоны). Это ограничение libdatrie, которое позволяет эффективно хранить автомат состояний и поддерживать юникод — сначала указываются допустимые диапазоны unicode-символов, потом unicode-ключи транслируются в более компактное внутреннее представление.

Мне кажется, что можно было бы на практике обойтись однобайтовыми кодировками вроде cp1251 и иметь примерно тот же функционал и эффективность, но подход с диапазонами символов тоже работает неплохо, уж сделано так в libdatrie и ладно. Я поэтому пишу, что «не поддерживает юникод — это ОК» — в случае с trie однобайтовая кодировка может быть удобным вариантом.

Потом с trie можно работать как со словарем:

>>> trie[u'foo'] = 5
>>> trie[u'foobar'] = 10
>>> trie[u'bar'] = 'bar value'
>>> trie.setdefault(u'foobar', 15)
10

>>> u'foo' in trie
True

>>> trie[u'foo']
5

Ключи должны быть юникодными :) В питоне 3.х это вполне естественно, в 2.x в примерах приходится ставить букву u, уж простите. Значения могут быть любыми питоньими объектами, что нетипично для C-реализаций trie (обычно там целые число в качестве значений). На самом деле «внутри» значения и правда представляют собой целые числа, но datrie.Trie использует их как индексы в массиве «настоящих» значений. Для тех, кому эта фича не нужна (например, значения не интересуют совсем), в datrie есть datrie.BaseTrie, который немного быстрее и умеет хранить только числа.

Немного о скорости. Все замеры проводил на trie.Trie с сотней тысяч уникальных русский и английских слов (50/50) и int-значениями «1», другие замеры (на миллионе url'ов) можно глянуть тут или (еще лучше) проводить самим на своих данных. Замеры скорости я проводил только для того, чтоб иметь какое-то общее представление о порядке величин, и чтоб отслеживать регрессии в библиотеке, так что принимайте их соответственно; весь исходный код и данные есть в репозитории. Кроме того, нигде не буду писать об асимптотической сложности различных операций, т.к. не исследовал ее. В теории должно быть как у trie из википедии (например, получение элемента или поиск по префиксу — O(m), где m — длина ключа), но детали реализации (как libdatrie, так и моей обертки) могут все менять; если кто-нибудь построит соответствующие графики, буду, конечно же, благодарен.

Операции получения элемента, проверки in, обновления элемента работают в среднем раза в 2-3 медленнее, чем у стандартного словаря (т.е. быстро, на моем буке это от 1 до 3 млн операций в секунду для вышеупомянутого trie). Исключение — это вставка нового значения в trie, она работает значительно медленнее (порядка 50тыс операций в секунду для того же trie с русскими и английскими словами). При этом trie на таких данных занимает значительно меньше места в оперативной памяти: 3-5M (в зависимости от интерпретатора) vs 20M+ у обычного словаря (замеры памяти проводил коряво и за конкретные цифры не ручаюсь).

Т.е. datrie.Trie можно использовать как замену dict, если есть большое количество не очень длинных строк (например, слов или url'ов), данные в основном используются в режиме read-only (или «update-only») и хочется сэкономить оперативную память память ценой 2-3 кратного снижения скорости доступа.

Чтоб этот псевдо-read-only не очень смущал, trie можно сохранять в файл и загружать из файла:

>>> trie.save('my.trie')
>>> trie2 = datrie.Trie.load('my.trie')


Еще одна особенность trie — это то, что некоторые операции, которые в dict (и любой другой хэш-таблице) потребовали бы полного перебора или большого объема памяти для построения дополнительных индексов, в префиксном дереве работают почти так же быстро (а некоторые даже быстрее), как и простое получение значения.

Можно проверить, есть ли в trie элементы, ключи у которых начинаются с данного префикса (это даже быстрее, чем простая проверка in):

>>> trie.has_keys_with_prefix(u'fo')
True


Можно найти все префиксы данной строки, которые есть в данном trie (это медленнее, по тестам — 500-600тыс операций/сек):

>>> trie.prefixes(u'foobarbaz')
[u'foo', u'foobar']

>>> trie.prefix_items(u'foobarbaz')
[(u'foo', 5), (u'foobar', 10)]


Можно найти все элементы, у которых ключи начинаются с данной строки:

>>> trie.keys(u'fo')
[u'foo', u'foobar']

>>> trie.items(u'fo')
[(u'foo', 5), (u'foobar', 10)]

>>> trie.values(u'foob')
[10]


В последнем примере основное время сейчас тратится на построение результатов, а не на поиск, это можно оптимизировать; сейчас скорость была где-то порядка 150-300 тыс. значений/сек (например, при префиксе длиной 7 и в среднем 3 значениями это 70тыс операций/сек).

У datrie.Trie есть еще различные методы, справку по ним и дополнительные результаты замеров скорости можно посмотреть в README в репозитории.

Под pypy у меня все завелось на дебиане (раз в 10 медленнее, чем под cpython); на маке под pypy не завелось (с обычным питоном должно работать и на маке, и на линуксе, и под windows). C API — расширения в pypy всегда будут медленные, т.к. работают через костыль cpyext, а cython генерирует именно C API расширения. Можно было написать обертку на ctypes, но она была бы медленной под обычным питоном (и не факт что быстрой под pypy) + распространять ctypes-расширения неудобно. Ребята из pypy пилят сейчас cffi, обещают, что он быстрый будет (и под cpython, и под pypy). Плюс, возможно, cython научится когда-нибудь генерировать не C API расширения, а cffi-расширения. Вот тогда, возможно, и наступит счастье) А пока что с pypy делать — не знаю. Ничего не буду, наверное; как-то работает все под линуксом и ладно.

В процессе реализации столкнулся с тормознутым кодеком utf_32_le в питоне. На это есть баг с патчем ( bugs.python.org/issue15027 ), но патч еще не закоммитили. Изначально все операции в datrie работали в 10 раз медленнее, но потом удалось в одном месте обойтись без кодирования строки стандартным питоньим utf_32_le и все заработало получше. Этот кодек используется еще в паре «горячих» мест, так что, думаю, если его ускорят, некоторые операции в datrie могут заработать еще до 2 раз быстрее.

Итерация по дереву сейчас тоже не самая эффективная, это связано с особенностями интерфейса libdatrie. Но автор libdatrie — отличный чел и собирается все поправить, так что перспективы неплохие.

Как обычно, патчи, баг-репорты, идеи, бенчмарки, пулл-реквесты и тд приветствуются!

github / bitbucket, кому что удобнее.
Коробов Михаил @kmike
карма
334,7
рейтинг 0,0
Пользователь
Реклама помогает поддерживать и развивать наши сервисы

Подробнее
Спецпроект

Самое читаемое Разработка

Комментарии (18)

  • +3
    Спасибо!
  • 0
    Спасибо большое! Очень полезная библиотека и её хорошее описание.
  • 0
    А для чего нужны отдельные save и load, чем не устраивает pickle/cPickle dump/load?
    • +2
      Тем, что для реализации pickle-протокола для datrie.Trie пришлось бы писать свою сериализацию trie в байты (на С или Cython).

      Не то чтобы это было очень сложно, но работы хватит, за час не напишу, наверное. В libdatrie уже есть поддержка сохранения trie в файл на диске, поэтому в datrie реализовано сохранение в файл: методы save и write. Метод write не документирован, но позволяет писать trie в открытый питоний file.

      Поддержка pickle была бы хорошей фичей, но ее пока нет, т.к. ее делать было сложнее, чем просто сохранение в файл. Патчи, понятное дело, приветствуются)
      • +2
        Можно «схитрить» (не бейте сильно...):

        import cPickle, datrie, string
        from datrie import AlphaMap
        
        class PickledTrie(datrie.Trie):
            _alpha_map = None
        
            def __init__(self, alpha_map=None, _create=True):
                self._alpha_map = alpha_map
                super(PickledTrie, self).__init__(AlphaMap(alpha_map), _create)
        
            def __reduce__(self):
                return PickledTrie, (self._alpha_map, ), None, None, iter(self.items())
            
            def __setitem__(self, k, v):
                super(PickledTrie, self).__setitem__(k, v)
        
        
        pt = PickledTrie(string.ascii_lowercase)
        pt[u'foo'] = u'bar'
        
        with open('tmp.dat', 'wb') as f:
            cPickle.dump(pt, f)
        
        with open('tmp.dat', 'rb') as f:
            z = cPickle.load(f)
            print z.items()
        
        
        • 0
          __setitem__ лишний, забыл стереть…
        • 0
          Хм, да, можно. Но это такой «хак», понятное дело)

          Метод .items() без префиксов сейчас тормозной (раз в 20-30 медленнее, чем у словаря); для сохранения и загрузки придется дублировать trie в памяти в «развернутом» виде (а с текущими save-load можно большой список слов в trie загрузить без загрузки всех отдельных слов как питоньих строк); у AlphaMap есть еще параметр ranges, который может быть использовать удобнее, чем alphabet (и который тут никак не передать).

          Но мне нравится) Необходимость создавать отдельный AlphaMap хорошо бы все равно из API убирать.

          Вот только если потом пиклинг переписать по-нормальному, то станет нельзя загружать старые trie, сохраненные этим вариантом. Способы обхода есть (например, распиклить trie старой версией, сохранить в файл через save и загрузить через load новой), но мне пока неочевидно, что игра стоит свеч, т.к. не очень представляю, зачем именно пиклинг может понадобиться. А правда, зачем?
          • 0
            А фиг его знает, кому «чистый» пиклинг нужен, мне хотелось поигратся… Хотя да, кривовато вышло.
            • 0
              Я, в принципе, за кривые решения, если они позволяют что-то сделать полезное и их потом можно заменить на «прямые» без смены интерфейса) Но тут вылезает досадная неприятность с обновлением со старой версии + польза не очевидна, поэтому в библиотеку этот вариант добавлять не хочется.

              Можно в вики добавить, вдруг кому пригодится)
              • 0
                Я не против… Кстати, с удовольствием попробую решить другие проблемы/накатать фичи для datrie. Есть предложения?
                • 0
                  Добавил в вики. Предложения — конечно, есть :)

                  * реализовать больше методов стандартного питоньего словаря (простор есть — от простых get и pop до хитрых viewitems);

                  * подумать, как убрать AlphaMap из публичного интерфейса (и datrie.new, наверное, тоже);

                  * добавить поддержку ключей-байтовых строк для второго питона (обязательно проверив на бенчмарках, что это ничего не затормозило для юникодных строк) — скорее всего, просто декодировать их в юникод из ascii; если затормозило — не добавлять, а явно задокументировать «не будет этого», отрицательный результат — тоже результат;

                  * реализовать аналог lookup c deepsearch='choose' из Trie::Trie. В datrie это, видимо, выльется в 3 метода first_key(prefix), first_item(prefix), first_value(prefix). С другой стороны, эти методы могут быть не очень нужны, если iteritems(prefix) будет реализован, т.к. они станут просто чуть более эффективной версией next(trie.iteritems(prefix)). С третьей стороны, они все равно могут быть полезны, т.к. они будут более эффективны + еще неизвестно, когда iteritems получится сделать.

                  * скрипт для бенчмарков довольно неудобный — все запускается сразу и ждать долговато; какой-то интерфейс командной строки к нему был бы очень кстати;

                  * любые дополнительные бенчмарки — это хорошо (и не только скорости: например, нормальный скрипт замера оперативной памяти — см. psutil и github.com/fabianp/memory_profiler ); вдруг что-то интересно узнать самому?;

                  * можно в самой libdatrie реализовать итерацию по дереву с каким-то iterator-object, и функциями в духе «next_state(const Trie* trie, IteratorState* state)» и послать патч автору libdatrie; после этого можно переписывать BaseTrie._walk_prefixes и добавлять публичный питоний интерфейс для гуляния по дереву. Theppitak Karoonboonyanan собирался в libdatrie поддержку итерации сам добавить, но не факт, что скоро;

                  * возможно, __len__ у Trie может быть быстрее (т.к. в отличие от BaseTrie у нас есть self.values — но как это будет работать совместно с __delitem__?)

                  * попробовать прогнать тесты на narrow-сборке питона (там сейчас довольно ковбойский метод кодирования питоньего юникода в AlphaChar*, в портабельности которого я до конца не уверен);

                  * ну или просто использовать библиотеку для своих нужд — недостающие фичи или проблемы сами вылезут со временем)
                  • 0
                    Вечером пошурую немного, может толк выйдет. Кстати, спасибо за extension!
        • 0
          Тот вариант, что родился у меня вчера в голове… Что-то в таком роде (не проверял, код из головы, но идея, думаю, ясна):
          import cPickle, datrie, tempfile, os
          
          class PicklableTrie(datrie.Trie):
              def __reduce__(self):
                  tempHandle, tempName = tempfile.mkstemp()
                  try:
                      os.close(tempHandle) # close the file we opened; now we'd save into it
                      self.save(tempName)
                      with open(tempName, 'rb') as inp:
                          return datrie.__version__, inp.read()
                  finally:
                      os.unlink(tempName)
          
              def __setstate__(self, newState):
                  assert datrie.__version__ == newState[0]
                  tempHandle, tempName = tempfile.mkstemp()
                  try:
                      os.close(tempHandle) # close the file we opened; now we'd save into it
                      with open(tempName, 'wb') as out:
                          out.write(newState[1])
                      self.load(tempName)
                  finally:
                      os.unlink(tempName)

          А вообще всё было бы куда проще и элегантнее, если бы save/load принимали не строки-имена файлов, а file-like object'ы. Тогда этот tempfile не нужен, и pickle обеспечивается за счёт класса StringIO.StringIO.
          • 0
            Кстати, если я правильно понимаю файлики-сорцы, то нужные для использования file-like object'a пункты можно довольно легко портировать из самой либы в обёртку, либо заменить fileutils.c, научив функции file_write_* и file_read_* работать с file-like питоновскими объектами.
            • 0
              Угу, про костыль с временными файлами тоже думал. В принципе, этот вариант даже получше, чем явная сериализация items, т.к. требует меньше оперативной памяти.

              Понятное дело, было бы проще, если б file-like объекты все поддерживало) Но там одним изменением fileutils не отделаешься, т.к. вся библиотека работает с FILE* и подразумевает наличие «настоящего» файла (tail.c/tail_fread, darray.c/da_fread).

              Кроме того, хорошо бы отвязываться от апстрима по минимуму — если можно в C-код изменений не вносить, то не вносить их; если и вносить, то такие, которые можно в апстрим пропихнуть (== неспецифичные для питона) — например, фикс для MSVC2008 так в транке libdatrie появился недавно.

              В том, чтоб встать на путь правки самой библиотеки, есть плюсы — например, можно было бы попробовать минимизировать преобразования PyUnicode <-> AlphaChar* + возможно, сделать TrieData == PyObject*. Кстати, ruby-расширения для libdatrie на этот путь и встало. Но и минусы очевидны — меньше пользы для мировой цивилизации, сложнее обновляться, больше кода поддерживать.

              Поэтому конкретно для «правильной» сериализации я бы предпочел дополнительные функции (на С или Cython), а не правку библиотеки.
              • 0
                Ну тогда можно костыль со временными файлами загнать прямо в Cython-обёртку: www.cplusplus.com/reference/clibrary/cstdio/tmpfile/ и *_fwrite() из libdatrie
                • 0
                  Можно. Нет желания сделать pull request?) Или сам могу добавить.

                  Мне вот только гугл говорит, что tmpfile из stdio может требовать административных прав под windows, так что я бы не стал его использовать; может быть лучше использовать менее прямой вариант с питоним tempfile.
  • +2
    Отличная статья, автору спасибо.
    Надо бы допилить до вменяемого состояния свою .NET реализацию по мотивам которой я и писал ту статью по ссылке. Интересно будет сравнить быстродействие, хотя вряд ли оно приблизится к сишной реализации.

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