<commit_before>import itertools
from BTrees.Length import Length
from zope.index.text import TextIndex as ZopeTextIndex
from zope.index.text.lexicon import CaseNormalizer
from zope.index.text.lexicon import Lexicon as _Lexicon
from zope.index.text.lexicon import Splitter
from zope.index.text.lexicon import _text2list
from zope.index.text.lexicon import StopWordRemover
from zope.index.text.okapiindex import OkapiIndex as _OkapiIndex
from repoze.catalog.indexes.text import CatalogTextIndex as _CatalogTextIndex
from zerodb import trees
from zerodb.catalog.indexes.common import CallableDiscriminatorMixin
from zerodb.storage import prefetch, prefetch_trees, parallel_traversal
from zerodb.catalog.indexes.pwid import PersistentWid


class Lexicon(_Lexicon):
    family = trees.family32  # In comparison with standard Lexicon, use bigger buckets

    def __init__(self, *pipeline):
        self._wids = self.family.OI.BTree()
        self._words = self.family.IO.BTree()
        self.wordCount = Length()
        self._pipeline = pipeline

    def sourceToWordIds(self, text):
        if text is None:
            text = ''
        last = _text2list(text)
        for element in self._pipeline:
            last = element.process(last)
        if not isinstance(self.wordCount, Length):
            # Make sure wordCount is overridden with a BTrees.Length.Length
            self.wordCount = Length(self.wordCount())
        # Strategically unload the length value so that we get the most
        # recent value written to the database to minimize conflicting wids
        # Because length is independent, this will load the most
        # recent value stored, regardless of whether MVCC is enabled
        self.wordCount._p_deactivate()
        parallel_traversal(self._wids, last)
        return list(map(self._getWordIdCreate, last))


class OkapiIndex(_OkapiIndex):

    def clear(self):
        # wid -> {docid -> weight}; t -> D -> w(D, t)
        self._wordinfo = trees.family32.IO.BTree()
        # XXX
        # Scalability of this Zope's approach is pretty bad (esp. when many documents).
        # Following is needed:
        # _wordinfo = BTree of (word, weight, docid) - TreeSet could be used instead
        # searching a keyword will be as _wordinfo.keys((word_start, None, None)),
        # already sorted by weight (just has to be multiplied by idf)
        # this works for both search and glob_search
        # However, when searching multiple keywords,
        # Need to find an efficient (logarithmic) algorithm of
        # incremental, weighted set intersection
        # Even without efficient intersection, it is faster and more secure anyway
        # XXX

        # docid -> weight
        self._docweight = self.family.IF.BTree()

        # docid -> WidCode'd list of wids (~1/4 of document size)
        # Used for un-indexing, and for phrase search.
        self._docwords = self.family.IO.BTree()

        # Use a BTree length for efficient length computation w/o conflicts
        self.wordCount = Length()
        self.documentCount = Length()

    def _mass_add_wordinfo(self, wid2weight, docid):
        dicttype = type({})
        # self._wordinfo - IOBTree of docid -> weight trees
        get_doc2score = self._wordinfo.get
        new_word_count = 0

        # Fill up cache for performance over the network
        wids = wid2weight.keys()
        parallel_traversal(self._wordinfo, wids)
        parallel_traversal(map(get_doc2score, wids), [docid] * len(wids))

        from time import time
        for wid, weight in wid2weight.items():
            doc2score = get_doc2score(wid)
            if doc2score is None:
                doc2score = {}
                new_word_count += 1
            elif (isinstance(doc2score, dicttype) and
                  len(doc2score) == self.DICT_CUTOFF):
                doc2score = self.family.IF.BTree(doc2score)
            doc2score[docid] = weight
            self._wordinfo[wid] = doc2score  # not redundant:  Persistency!
        try:
            self.wordCount.change(new_word_count)
        except AttributeError:
            # upgrade wordCount to Length object
            self.wordCount = Length(len(self._wordinfo))

    def index_doc(self, docid, text):
        if docid in self._docwords:
            return self._reindex_doc(docid, text)
        wids = self._lexicon.sourceToWordIds(text)
        wid2weight, docweight = self._get_frequencies(wids)
        self._mass_add_wordinfo(wid2weight, docid)
        self._docweight[docid] = docweight
        self._docwords[docid] = PersistentWid.encode_wid(wids)
        try:
            self.documentCount.change(1)
        except AttributeError:
            # upgrade documentCount to Length object
            self.documentCount = Length.Length(len(self._docweight))
        count = len(wids)
        self._change_doc_len(count)
        return count

    def _reindex_doc(self, docid, text):
        # Touch as few docid->w(docid, score) maps in ._wordinfo as possible.
        self._change_doc_len(-self._docweight[docid])

        old_wids = self.get_words(docid)
        old_wid2w, old_docw = self._get_frequencies(old_wids)

        new_wids = self._lexicon.sourceToWordIds(text)
        new_wid2w, new_docw = self._get_frequencies(new_wids)

        old_widset = self.family.IF.TreeSet(old_wid2w.keys())
        new_widset = self.family.IF.TreeSet(new_wid2w.keys())

        IF = self.family.IF
        in_both_widset = IF.intersection(old_widset, new_widset)
        only_old_widset = IF.difference(old_widset, in_both_widset)
        only_new_widset = IF.difference(new_widset, in_both_widset)
        del old_widset, new_widset

        for wid in only_old_widset.keys():
            self._del_wordinfo(wid, docid)

        for wid in only_new_widset.keys():
            self._add_wordinfo(wid, new_wid2w[wid], docid)

        for wid in in_both_widset.keys():
            # For the Okapi indexer, the "if" will trigger only for words
            # whose counts have changed.  For the cosine indexer, the "if"
            # may trigger for every wid, since W(d) probably changed and
            # W(d) is divided into every score.
            newscore = new_wid2w[wid]
            if old_wid2w[wid] != newscore:
                self._add_wordinfo(wid, newscore, docid)

        self._docweight[docid] = new_docw
        self._docwords[docid] = PersistentWid.encode_wid(new_wids)
        return len(new_wids)

    def get_words(self, docid):
        """Return a list of the wordids for a given docid."""
        return self._docwords[docid].decode_wid()

    def _search_wids(self, wids):
        # Bulk-fetch all the info we want to use
        if len(wids) > 1:
            parallel_traversal(self._wordinfo, wids)
        prefetch_trees([self._wordinfo[wid] for wid in wids])

        docids = list(set(itertools.chain(
            *[self._wordinfo[wid].keys() for wid in wids])))
        if len(docids) > 1:
            parallel_traversal(self._docweight, docids)

        return super(OkapiIndex, self)._search_wids(wids)


class IncrementalLuceneIndex(object):
    """
    Using Lucene's practical scoring function and incremental search
    """
    # Using Okapi index requires us to fetch all docids for all word ids.
    # That makes scalability of the search bad, even for a single keyword
    # This class uses Lucene's practical scoring function:
    # https://www.elastic.co/guide/en/elasticsearch/guide/current/practical-scoring-function.html

    # In brief, practical scoring is this:
    #   t - keyword (or term) which is searched
    #   D - document
    #   TF(t, D) = sqrt(f(t, D))
    #   IDF(t) = 1 + log(N_docs / (N_docs_for_term + 1))
    #   score(term, D) = TF * IDF**2 / sqrt(N_terms_in_D)
    #   score(D) = sum(score(term, D)) / sqrt(sum(IDF(t)**2))

    # When we search for one keyword, we can keep document UIDs
    # Ordered by w_t = TF(t, D) / sqrt(N_terms_in_D),
    # for example in a TreeSet((-w_t, uid)).
    # Reading it incrementally gives docs ordered from most to least relevant.
    pass


class CatalogTextIndex(CallableDiscriminatorMixin, _CatalogTextIndex):
    family = trees.family32

    def __init__(self, discriminator, lexicon=None, index=None):
        self._init_discriminator(discriminator)

        self._not_indexed = self.family.IF.Set()

        lexicon = lexicon or Lexicon(Splitter(), CaseNormalizer(), StopWordRemover())
        index = index or OkapiIndex(lexicon, family=self.family)

        ZopeTextIndex.__init__(self, lexicon, index)
        self.clear()<commit_msg>How to do fulltext search incrementally<commit_after>import itertools
from BTrees.Length import Length
from zope.index.text import TextIndex as ZopeTextIndex
from zope.index.text.lexicon import CaseNormalizer
from zope.index.text.lexicon import Lexicon as _Lexicon
from zope.index.text.lexicon import Splitter
from zope.index.text.lexicon import _text2list
from zope.index.text.lexicon import StopWordRemover
from zope.index.text.okapiindex import OkapiIndex as _OkapiIndex
from repoze.catalog.indexes.text import CatalogTextIndex as _CatalogTextIndex
from zerodb import trees
from zerodb.catalog.indexes.common import CallableDiscriminatorMixin
from zerodb.storage import prefetch, prefetch_trees, parallel_traversal
from zerodb.catalog.indexes.pwid import PersistentWid


class Lexicon(_Lexicon):
    family = trees.family32  # In comparison with standard Lexicon, use bigger buckets

    def __init__(self, *pipeline):
        self._wids = self.family.OI.BTree()
        self._words = self.family.IO.BTree()
        self.wordCount = Length()
        self._pipeline = pipeline

    def sourceToWordIds(self, text):
        if text is None:
            text = ''
        last = _text2list(text)
        for element in self._pipeline:
            last = element.process(last)
        if not isinstance(self.wordCount, Length):
            # Make sure wordCount is overridden with a BTrees.Length.Length
            self.wordCount = Length(self.wordCount())
        # Strategically unload the length value so that we get the most
        # recent value written to the database to minimize conflicting wids
        # Because length is independent, this will load the most
        # recent value stored, regardless of whether MVCC is enabled
        self.wordCount._p_deactivate()
        parallel_traversal(self._wids, last)
        return list(map(self._getWordIdCreate, last))


class OkapiIndex(_OkapiIndex):

    def clear(self):
        # wid -> {docid -> weight}; t -> D -> w(D, t)
        self._wordinfo = trees.family32.IO.BTree()
        # XXX
        # Scalability of this Zope's approach is pretty bad (esp. when many documents).
        # Following is needed:
        # _wordinfo = BTree of (word, weight, docid) - TreeSet could be used instead
        # searching a keyword will be as _wordinfo.keys((word_start, None, None)),
        # already sorted by weight (just has to be multiplied by idf)
        # this works for both search and glob_search
        # However, when searching multiple keywords,
        # Need to find an efficient (logarithmic) algorithm of
        # incremental, weighted set intersection
        # Even without efficient intersection, it is faster and more secure anyway
        # XXX

        # docid -> weight
        self._docweight = self.family.IF.BTree()

        # docid -> WidCode'd list of wids (~1/4 of document size)
        # Used for un-indexing, and for phrase search.
        self._docwords = self.family.IO.BTree()

        # Use a BTree length for efficient length computation w/o conflicts
        self.wordCount = Length()
        self.documentCount = Length()

    def _mass_add_wordinfo(self, wid2weight, docid):
        dicttype = type({})
        # self._wordinfo - IOBTree of docid -> weight trees
        get_doc2score = self._wordinfo.get
        new_word_count = 0

        # Fill up cache for performance over the network
        wids = wid2weight.keys()
        parallel_traversal(self._wordinfo, wids)
        parallel_traversal(map(get_doc2score, wids), [docid] * len(wids))

        from time import time
        for wid, weight in wid2weight.items():
            doc2score = get_doc2score(wid)
            if doc2score is None:
                doc2score = {}
                new_word_count += 1
            elif (isinstance(doc2score, dicttype) and
                  len(doc2score) == self.DICT_CUTOFF):
                doc2score = self.family.IF.BTree(doc2score)
            doc2score[docid] = weight
            self._wordinfo[wid] = doc2score  # not redundant:  Persistency!
        try:
            self.wordCount.change(new_word_count)
        except AttributeError:
            # upgrade wordCount to Length object
            self.wordCount = Length(len(self._wordinfo))

    def index_doc(self, docid, text):
        if docid in self._docwords:
            return self._reindex_doc(docid, text)
        wids = self._lexicon.sourceToWordIds(text)
        wid2weight, docweight = self._get_frequencies(wids)
        self._mass_add_wordinfo(wid2weight, docid)
        self._docweight[docid] = docweight
        self._docwords[docid] = PersistentWid.encode_wid(wids)
        try:
            self.documentCount.change(1)
        except AttributeError:
            # upgrade documentCount to Length object
            self.documentCount = Length.Length(len(self._docweight))
        count = len(wids)
        self._change_doc_len(count)
        return count

    def _reindex_doc(self, docid, text):
        # Touch as few docid->w(docid, score) maps in ._wordinfo as possible.
        self._change_doc_len(-self._docweight[docid])

        old_wids = self.get_words(docid)
        old_wid2w, old_docw = self._get_frequencies(old_wids)

        new_wids = self._lexicon.sourceToWordIds(text)
        new_wid2w, new_docw = self._get_frequencies(new_wids)

        old_widset = self.family.IF.TreeSet(old_wid2w.keys())
        new_widset = self.family.IF.TreeSet(new_wid2w.keys())

        IF = self.family.IF
        in_both_widset = IF.intersection(old_widset, new_widset)
        only_old_widset = IF.difference(old_widset, in_both_widset)
        only_new_widset = IF.difference(new_widset, in_both_widset)
        del old_widset, new_widset

        for wid in only_old_widset.keys():
            self._del_wordinfo(wid, docid)

        for wid in only_new_widset.keys():
            self._add_wordinfo(wid, new_wid2w[wid], docid)

        for wid in in_both_widset.keys():
            # For the Okapi indexer, the "if" will trigger only for words
            # whose counts have changed.  For the cosine indexer, the "if"
            # may trigger for every wid, since W(d) probably changed and
            # W(d) is divided into every score.
            newscore = new_wid2w[wid]
            if old_wid2w[wid] != newscore:
                self._add_wordinfo(wid, newscore, docid)

        self._docweight[docid] = new_docw
        self._docwords[docid] = PersistentWid.encode_wid(new_wids)
        return len(new_wids)

    def get_words(self, docid):
        """Return a list of the wordids for a given docid."""
        return self._docwords[docid].decode_wid()

    def _search_wids(self, wids):
        # Bulk-fetch all the info we want to use
        if len(wids) > 1:
            parallel_traversal(self._wordinfo, wids)
        prefetch_trees([self._wordinfo[wid] for wid in wids])

        docids = list(set(itertools.chain(
            *[self._wordinfo[wid].keys() for wid in wids])))
        if len(docids) > 1:
            parallel_traversal(self._docweight, docids)

        return super(OkapiIndex, self)._search_wids(wids)


class IncrementalLuceneIndex(object):
    """
    Using Lucene's practical scoring function and incremental search
    """
    # Using Okapi index requires us to fetch all docids for all word ids.
    # That makes scalability of the search bad, even for a single keyword
    # This class uses Lucene's practical scoring function:
    # https://www.elastic.co/guide/en/elasticsearch/guide/current/practical-scoring-function.html

    # In brief, practical scoring is this:
    #   t - keyword (or term) which is searched
    #   D - document
    #   TF(t, D) = sqrt(f(t, D))
    #   IDF(t) = 1 + log(N_docs / (N_docs_for_term + 1))
    #   score(term, D) = TF * IDF**2 / sqrt(N_terms_in_D)
    #   score(D) = sum(score(term, D)) / sqrt(sum(IDF(t)**2))

    # When we search for one keyword, we can keep document UIDs
    # Ordered by w_t = TF(t, D) / sqrt(N_terms_in_D),
    # for example in a TreeSet((-w_t, uid)).
    # Reading it incrementally gives docs ordered from most to least relevant.

    # Now the harder part (not in Lucene)

    # When we want to do a multi-keyword search, it is, in a way, weighted union
    # of searches for all keywords.
    # E.g. now:
    #   w(D) = sum(w_t * IDF(t)**2)   // no need to normalize by sqrt(sum(idf**2))
    # weight calculation and sort in memory, which fetches docweights for all terms again.
    # This is slow when e2e encrypted and fetched remotely, and also reveals sum(N_docs_for_term) to the server

    # Here's how we do the same incrementally ("lazy")
    # We allow elements to be shifted by max_out_of_order = 3 from their ideal position.
    # We read beginnings of sorted lists of (w_i, uid_i) for each word.
    # After we calculate the sum,
    # w_j_known = sum(w_j * IDF(t_j)**2, t: all terms for which we know w_j for uid_i)

    # for i_th element the weight could be between:
    #   w_min_i(t_j) = sum(w_min(t_j) * IDF(t_j)**2: all terms for which we don't know w_j) + w_i_known
    #   and
    #   w_max_i(t_j) = sum(w_max(t_j) * IDF(t_j)**2: all terms for which we don't know w_j) + w_i_known

    # For each term, we have current window of docs ordered by weight which still haven't been returned.
    # It is likely that some term is particularly selective and largely outweights others,
    # E.g. for almost all elements in the current window, weight ranges (w_min_i, w_max_i) don't overlap
    # for different i. Or, perhaps, elements are misplaced by no more than max_out_of_order positions.

    # We choose a term for which all positions at the tail of the window can never appear within the limit # of docs.
    # If there are none, we expand most promising wid set by factor of 2 until we get this situation.
    # Out of those terms (if there are multiple), we select one which has smallest number of
    # doc uids misplaced.

    # For misplaced docids, we fetch their weights. Now we know order within the specified
    # number of documents.

    # For this to work, we need to keep
    # BTree(wids -> (TreeSet(w_i, uid_i), BTree(uid_i -> w_i), Length))
    # Length is for calculating IDF


class CatalogTextIndex(CallableDiscriminatorMixin, _CatalogTextIndex):
    family = trees.family32

    def __init__(self, discriminator, lexicon=None, index=None):
        self._init_discriminator(discriminator)

        self._not_indexed = self.family.IF.Set()

        lexicon = lexicon or Lexicon(Splitter(), CaseNormalizer(), StopWordRemover())
        index = index or OkapiIndex(lexicon, family=self.family)

        ZopeTextIndex.__init__(self, lexicon, index)
        self.clear()
{'fraction_non_alphanumeric': 0.08596181780452376, 'fraction_numerical': 0.004876530400498028, 'mean_word_length': 2.698580199539524, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 25, 'https://': 2, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '794112', 'n_tokens_mistral': 5998, 'n_tokens_neox': 5783, 'n_words': 1727}