sqlite-fts-userspace.py

import binascii
import sqlite3
import collections
import math
import re

RE_SPLIT = re.compile(r'[^A-Za-z0-9_]+')

def tokenize(s):
    # this is super basic and should be replaced with a real tokenizer
    return [word[:-2] for word in RE_SPLIT.split(s.lower()) if len(word) > 2]

conn = sqlite3.connect(':memory:')
c = conn.cursor()

c.execute('''
create table terms (
    search_index_id bigint not null,
    term_id bigint not null,
    doc_id bigint not null,
    tf double not null,
    primary key (term_id, doc_id)
)
''')

c.execute('''
create index idx_terms on terms (search_index_id, term_id, tf)
''')

def add_document(cursor, search_index_id, doc_id, tokens):
    counter = collections.Counter(tokens)
    rows = ((search_index_id, binascii.crc32(bytes(key, 'utf8')), doc_id, math.sqrt(count)) for key, count in counter.items())
    cursor.executemany('insert into terms (search_index_id, term_id, doc_id, tf) values (?, ?, ?, ?)', rows)

def search(cursor, search_index_id, tokens, n_candidates = 10000):
    # implementation of https://www.elastic.co/guide/en/elasticsearch/guide/current/practical-scoring-function.html
    token_ids = set([binascii.crc32(bytes(key, 'utf8')) for key in tokens])
    
    # first get idf for each term
    cursor.execute('select count(*) from terms where search_index_id=?', (search_index_id,))
    num_docs, = cursor.fetchone()
    cursor.execute('select term_id, count(*) from terms where term_id in (%s) group by term_id' % (','.join(str(x) for x in token_ids),))
    idf = {}
    for term_id, doc_freq in cursor.fetchall():
        idf[term_id] = 1 + math.log(num_docs / (doc_freq + 1))
        
    # now get the most likely candidate doc_ids
    limit = n_candidates * len(token_ids)
    cursor.execute('select term_id, tf, doc_id from terms where search_index_id=%d and term_id in (%s) order by tf desc limit %d' % (search_index_id, ','.join(str(x) for x in token_ids), limit))
    
    # score each candidate doc_id using tfidf * coord
    results = collections.defaultdict(lambda: 0)
    coords = collections.defaultdict(lambda: 0)
    for term_id, tf, doc_id in cursor.fetchall():
        results[doc_id] += tf * idf[term_id]
        coords[doc_id] += 1
        
    # calc coord (percent of terms in the query that match the doc)
    for doc_id, num_matched in coords.items():
        results[doc_id] *= num_matched / len(token_ids)
        
    return sorted(((score, doc_id) for (doc_id, score) in results.items()), reverse=True)