Pythonで全文検索を実装してみた

JavaScript でやってるのを見かけたので。

 cf. JavaScriptで全文検索(N-gram)を実装してみる! – Simple is Beautiful.

アルゴリズムは N-gram っていう方法のうちでも簡単な uni-gram っていうみたい。詳しくはリンク先の記事を見て。

記事の説明を読んで、なんとなく理解したので Python で書いてみた。リンク先の JavaScript の実装はファイル名を参考にしたくらいで、ちゃんと読んでない。

できたのはこんな感じ。

  • create_index.py でインデックスを作って、search.py で検索
  • 検索対象のファイルはテキストファイルのみ。documents ディレクトリに入ってる
  • インデックスファイルは indexes ディレクトリ内につくる
  • document ID とファイル名の対応は docs.json ファイル

インデックスを作る create_index.py

from unigram import document
import json
import os


DOC_DIR = 'documents'
INDEX_DIR = 'indexes'
DOC_DATA = 'docs.json'


files = [os.path.join(DOC_DIR, f) for f in os.listdir(DOC_DIR)]
docs = {}
doc_id = 0
for file in files:
    with open(file, 'r') as f:
        text = f.read()
        tokens = document.tokenize(text)
        index = document.classify(tokens)
        document.save_index(index, doc_id, INDEX_DIR)
    docs[str(doc_id)] = {'name': os.path.basename(file), 'path': file}
    doc_id += 1

with open(DOC_DATA, 'w') as f:
    json.dump(docs, f, indent=2)

検索コマンド search.py

from unigram import document
import json
import os
import sys


INDEX_DIR = 'indexes'
DOC_DATA = 'docs.json'


string = sys.argv[1]

with open(DOC_DATA, 'r') as f:
    docs = json.load(f)
fcount = len(docs)

index_files = list(map(lambda x: os.path.join(INDEX_DIR, x), os.listdir(INDEX_DIR)))
index = {}
for file in index_files:
    c = chr(int(os.path.basename(file).replace('.index', '')))
    with open(file, 'r') as f:
        index[c] = document.parse_index(f.read())

m = list(map(lambda x: index[x], list(string)))

for i in range(fcount):
    doc_id = str(i)
    if not all(map(lambda x: doc_id in x.keys(), m)):
        continue
    n = list(map(lambda x: x[doc_id], m))
    s = set(n[0])
    for s1 in n[1:]:
        s = set(list(map(lambda x: x + 1, s)))
        s = s & set(s1)
    if len(s) > 0:
        pos = list(map(lambda x: x - len(string) + 1, s))
        pos.sort()
        print(docs[doc_id]['name'], pos)

両方から使うモジュール unigram/document.py

import os


def tokenize(text):
    return list(text)


def classify(token_list):
    tokens = {}
    pos = 0
    for t in token_list:
        if not t in tokens:
            tokens[t] = []
        tokens[t].append(pos)
        pos += 1
    return tokens


def save_index(index, doc_id, index_dir):
    for c, idx in index.items():
        l = str(doc_id) + ':' + ','.join(list(map(lambda x: str(x), idx))) + '\n'
        with open(os.path.join(index_dir, str(ord(c)) + '.index'), 'a') as f:
            f.write(l)


def parse_index(content):
    l = content.split('\n')
    l.pop()
    index = {}
    for l2 in l:
        a = l2.split(':')
        index[a[0]] = [int(x) for x in a[1].split(',')]
    return index

とにかくまずは動くものを、ってことで作ったので、コードが整理されてないのは目を瞑ってほしい(後で直す)。

検索対象ファイルのサンプルには、別のプロジェクトで書いた Ruby のソースファイル。

takatoh@apostrophe $ ls documents
Gemfile       Rakefile  boot.rb    config.yaml
Gemfile.lock  app.rb    config.ru  config.yaml.example

インデックスを作る。

takatoh@apostrophe $ python create_index.py

できたインデックスファイルがこれ。

takatoh@apostrophe $ ls indexes
10.index   111.index  123.index  44.index  58.index  71.index  84.index
100.index  112.index  124.index  45.index  60.index  72.index  85.index
101.index  113.index  125.index  46.index  61.index  73.index  87.index
102.index  114.index  126.index  47.index  62.index  74.index  89.index
103.index  115.index  32.index   48.index  63.index  75.index  91.index
104.index  116.index  34.index   49.index  64.index  76.index  93.index
105.index  117.index  35.index   50.index  65.index  77.index  95.index
106.index  118.index  36.index   51.index  66.index  78.index  97.index
107.index  119.index  39.index   52.index  67.index  79.index  98.index
108.index  120.index  40.index   53.index  68.index  80.index  99.index
109.index  121.index  41.index   55.index  69.index  82.index
110.index  122.index  43.index   56.index  70.index  83.index

「require」という文字列を検索してみる。ファイル名と出現位置(のリスト)が出力される。

takatoh@apostrophe $ python search.py require
boot.rb [0, 17]
Rakefile [0, 15, 32, 71]
config.ru [0]
app.rb [0, 23, 40, 56, 71]

大丈夫そうだ。