<commit_before>#===============================================================================
# Copyright 2007 <NAME>
# 
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# 
#    http://www.apache.org/licenses/LICENSE-2.0
# 
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#===============================================================================

"""
This module contains classes for writing to an index.
"""

from __future__ import with_statement
from array import array
from collections import defaultdict

from whoosh import index, postpool, reading, tables
from whoosh.fields import UnknownFieldError
from whoosh.util import fib

# Exceptions

class IndexingError(Exception):
    pass

# Constants

# TODO: it might be better style to have these classes actually
# implement the given merge policy, rather than just using them
# as constants.

class NO_MERGE(object):
    """Indicates the writer should NOT merge small segments upon completion."""
    def __init__(self): raise NotImplementedError

class MERGE_SMALL(object):
    """Indicates the writer should merge small segments upon completion."""
    def __init__(self): raise NotImplementedError

class OPTIMIZE(object):
    """Indicates the writer should merge ALL segments upon completion."""
    def __init__(self): raise NotImplementedError

# Writing classes

class IndexWriter(index.DeletionMixin):
    """High-level object for writing to an index. This object takes care of
    instantiating a SegmentWriter to create a new segment as you add documents,
    as well as merging existing segments (if necessary) when you finish.
    
    You can use this object as a context manager. If an exception is thrown
    from within the context it calls cancel(), otherwise it calls commit()
    when the context ends.
    """
    
    # This class is mostly a shell for SegmentWriter. It exists to handle
    # multiple SegmentWriters during merging/optimizing.
    
    def __init__(self, ix, term_blocksize = 1 * 1024, doc_blocksize = 8 * 1024,
                 vector_blocksize = 8 * 1024):
        """
        @param ix: the Index object you want to write to.
        @param blocksize: the block size for tables created by this writer.
        """
        
        # Obtain a lock
        self.locked = ix.lock()
        
        self.index = ix
        self.segments = ix.segments.copy()
        self.term_blocksize = term_blocksize
        self.doc_blocksize = doc_blocksize
        self.vector_blocksize = vector_blocksize
        self._segment_writer = None
        self._searcher = None
    
    def __enter__(self):
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type:
            self.cancel()
        else:
            self.commit()
    
    def _finish(self):
        self._segment_writer = None
        # Release the lock
        if self.locked:
            self.index.unlock()
    
    def segment_writer(self):
        """Returns the underlying SegmentWriter object."""
        
        if not self._segment_writer:
            self._segment_writer = SegmentWriter(self.index, self.term_blocksize,
                                                 self.doc_blocksize, self.vector_blocksize)
        return self._segment_writer
    
    def searcher(self):
        """Returns a searcher for the existing index."""
        if not self._searcher:
            self._searcher = self.index.searcher()
        return self._searcher
    
    def start_document(self):
        """Starts recording information for a new document. This should be followed by
        add_field() calls, and must be followed by an end_document() call.
        Alternatively you can use add_document() to add all fields at once.
        """
        self.segment_writer().start_document()
        
    def add_field(self, fieldname, text, stored_value = None):
        """Adds a the value of a field to the document opened with start_document().
        
        @param fieldname: The name of the field in which to index/store the text.
        @param text: The text to index.
        @type fieldname: string
        @type text: unicode
        """
        self.segment_writer().add_field(fieldname, text, stored_value = stored_value)
        
    def end_document(self):
        """
        Closes a document opened with start_document().
        """
        self.segment_writer().end_document()
    
    def add_document(self, **fields):
        """Adds all the fields of a document at once. This is an alternative to calling
        start_document(), add_field() [...], end_document().
        
        The keyword arguments map field names to the values to index/store.
        
        For fields that are both indexed and stored, you can specify an alternate
        value to store using a keyword argument in the form "_stored_<fieldname>".
        For example, if you have a field named "title" and you want to index the
        text "a b c" but store the text "e f g", use keyword arguments like this::
        
            add_document(title=u"a b c" _stored_title=u"e f g)
        """
        self.segment_writer().add_document(fields)
    
    def update_document(self, **fields):
        """Adds or replaces a document. At least one of the fields for which you
        supply values must be marked as 'unique' in the index's schema.
        
        The keyword arguments map field names to the values to index/store.
        
        For fields that are both indexed and stored, you can specify an alternate
        value to store using a keyword argument in the form "_stored_<fieldname>".
        For example, if you have a field named "title" and you want to index the
        text "a b c" but store the text "e f g", use keyword arguments like this::
        
            update_document(title=u"a b c" _stored_title=u"e f g)
        """
        
        # Check which of the supplied fields are unique
        unique_fields = [name for name, field
                         in self.index.schema.fields()
                         if name in fields and field.unique]
        if not unique_fields:
            raise IndexingError("None of the fields in %r are unique" % fields.keys())
        
        # Delete documents in which the supplied unique fields match
        searcher = self.searcher()
        for name in unique_fields:
            self.delete_by_term(name, fields[name], searcher = searcher)
        
        # Add the given fields
        self.add_document(**fields)
    
    def commit(self, mergetype = MERGE_SMALL):
        """Finishes writing and unlocks the index.
        
        @param mergetype: How to merge existing segments.
        @type mergetype: one of NO_MERGE, MERGE_SMALL, or OPTIMIZE
        """
        
        if self._segment_writer or mergetype is OPTIMIZE:
            self._merge_segments(mergetype)
        self.index.commit(self.segments)
        self._finish()
        
    def cancel(self):
        """Cancels any documents/deletions added by this object
        and unlocks the index.
        """
        self._finish()
    
    def _merge_segments(self, mergetype):
        if mergetype not in (NO_MERGE, MERGE_SMALL, OPTIMIZE):
            raise ValueError("Unknown merge type: %r" % mergetype)
        
        sw = self.segment_writer()
        
        segments = self.segments
        new_segments = index.SegmentSet()
        
        if mergetype is OPTIMIZE:
            # Merge all segments
            for seg in segments:
                sw.add_segment(self.index, seg)
        else:
            # Find sparse segments and merge them into the segment
            # currently being written.
            
            sorted_segment_list = sorted((s.doc_count_all(), s) for s in segments)
            total_docs = 0
            
            if mergetype is not NO_MERGE:
                # Merge sparse segments into the one we're
                # currently writing
                for i, (count, seg) in enumerate(sorted_segment_list):
                    if count > 0:
                        total_docs += count
                        if total_docs < fib(i + 5):
                            sw.add_segment(self.index, seg)
                        else:
                            new_segments.append(seg)
            else:
                new_segments = segments
        
        self._segment_writer.close()
        new_segments.append(sw.segment())
        self.segments = new_segments

# Constants for "special" fields
UNIQUE_COUNT = -2
TOTAL_COUNT = -1

class SegmentWriter(object):
    """
    Do not instantiate this object directly; it is created by the IndexWriter object.
    
    Handles the actual writing of new documents to the index: writes stored fields,
    handles the posting pool, and writes out the term index.
    """
    
    class DocumentState(object):
        def __init__(self):
            self.reset()
        
        def reset(self):
            #: Whether a document is currently in progress.
            self.active = False
            #: Maps field names to stored field contents for this document
            self.stored_fields = {}
            #: Keeps track of the last field that was added.
            self.prev_fieldnum = None
    
    def __init__(self, ix, term_blocksize, doc_blocksize, vector_blocksize, name = None):
        """
        @param ix: the Index object in which to write the new segment.
        @param name: the name of the segment.
        @param blocksize: the block size to use for tables created by this writer.
        """
        
        self.index = ix
        self.schema = ix.schema
        self.storage = ix.storage
        self.name = name or ix._next_segment_name()
        
        self.max_doc = 0
        self.term_count = 0
        self.max_weight = 0
        self.doc_field_lengths = defaultdict(list)
        self.field_length_totals = defaultdict(int)
        
        # Records the state of the writer wrt start_document/end_document.
        # None == not "in" a document.
        self._doc_state = SegmentWriter.DocumentState()
        self._scorable_fields = self.schema.scorable_fields()
        
        self.pool = postpool.PostingPool()
        
        # Create a temporary segment object just so we can access
        # its *_filename attributes (so if we want to change the
        # naming convention, we only have to do it in one place).
        tempseg = index.Segment(self.name, 0, 0, 0, None)
        
        # Open files for writing
        self.term_table = self.storage.create_table(tempseg.term_filename, postings = True,
                                                    blocksize = term_blocksize)
        self.doclength_table = self.storage.create_table(tempseg.doclen_filename,
                                                         blocksize = term_blocksize)
        self.docs_table = self.storage.create_table(tempseg.docs_filename,
                                                    blocksize = doc_blocksize, compressed = 9)
        
        self.vector_table = None
        if self.schema.has_vectored_fields():
            self.vector_table = self.storage.create_table(tempseg.vector_filename,
                                                          postings = True,
                                                          stringids = True,
                                                          blocksize = vector_blocksize)
            
    def segment(self):
        """Returns an index.Segment object for the segment being written."""
        return index.Segment(self.name, self.max_doc,
                             self.term_count, self.max_weight,
                             dict(self.field_length_totals))
    
    def close(self):
        """Finishes writing the segment (flushes the posting pool out to disk) and
        closes all open files.
        """
        
        if self._doc_state.active:
            raise IndexingError("Called SegmentWriter.close() with a document still opened")
        
        self._flush_pool()
        
        self.doclength_table.add_row((UNIQUE_COUNT), self.doc_field_lengths[UNIQUE_COUNT])
        self.doclength_table.add_row((TOTAL_COUNT), self.doc_field_lengths[TOTAL_COUNT])
        for fieldnum in self._scorable_fields:
            arr = array("i", self.doc_field_lengths[fieldnum])
            self.doclength_table.add_row((fieldnum), arr)
        self.doclength_table.close()
        <commit_msg>Changed NO_MERGE, MERGE_SMALL, and OPTIMIZE from constant classes into the actual implementations of the different merge policies.<commit_after>#===============================================================================
# Copyright 2007 <NAME>
# 
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# 
#    http://www.apache.org/licenses/LICENSE-2.0
# 
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#===============================================================================

"""
This module contains classes for writing to an index.
"""

from __future__ import with_statement
from array import array
from collections import defaultdict

from whoosh import index, postpool, reading, tables
from whoosh.fields import UnknownFieldError
from whoosh.util import fib

# Exceptions

class IndexingError(Exception):
    pass


# Merge policies

# A merge policy is a callable that takes the Index object,
# the SegmentWriter object, and the current SegmentSet
# (not including the segment being written), and returns an
# updated SegmentSet (not including the segment being
# written).

def NO_MERGE(ix, writer, segments):
    """This policy does not merge any existing segments.
    """
    return segments


def MERGE_SMALL(ix, writer, segments):
    """This policy merges small segments, where small is
    defined using a heuristic based on the fibonacci sequence.
    """
    
    newsegments = index.SegmentSet()
    sorted_segment_list = sorted((s.doc_count_all(), s) for s in segments)
    total_docs = 0
    for i, (count, seg) in enumerate(sorted_segment_list):
        if count > 0:
            total_docs += count
            if total_docs < fib(i + 5):
                writer.add_segment(ix, seg)
            else:
                newsegments.append(seg)
    return newsegments


def OPTIMIZE(ix, writer, segments):
    """This policy merges all existing segments.
    """
    for seg in segments:
        writer.add_segment(ix, seg)
    return index.SegmentSet()


# Writing classes

class IndexWriter(index.DeletionMixin):
    """High-level object for writing to an index. This object takes care of
    instantiating a SegmentWriter to create a new segment as you add documents,
    as well as merging existing segments (if necessary) when you finish.
    
    You can use this object as a context manager. If an exception is thrown
    from within the context it calls cancel(), otherwise it calls commit()
    when the context ends.
    """
    
    # This class is mostly a shell for SegmentWriter. It exists to handle
    # multiple SegmentWriters during merging/optimizing.
    
    def __init__(self, ix, term_blocksize = 1 * 1024, doc_blocksize = 8 * 1024,
                 vector_blocksize = 8 * 1024):
        """
        @param ix: the Index object you want to write to.
        @param blocksize: the block size for tables created by this writer.
        """
        
        # Obtain a lock
        self.locked = ix.lock()
        
        self.index = ix
        self.segments = ix.segments.copy()
        self.term_blocksize = term_blocksize
        self.doc_blocksize = doc_blocksize
        self.vector_blocksize = vector_blocksize
        self._segment_writer = None
        self._searcher = None
    
    def __enter__(self):
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type:
            self.cancel()
        else:
            self.commit()
    
    def _finish(self):
        self._segment_writer = None
        # Release the lock
        if self.locked:
            self.index.unlock()
    
    def segment_writer(self):
        """Returns the underlying SegmentWriter object."""
        
        if not self._segment_writer:
            self._segment_writer = SegmentWriter(self.index, self.term_blocksize,
                                                 self.doc_blocksize, self.vector_blocksize)
        return self._segment_writer
    
    def searcher(self):
        """Returns a searcher for the existing index."""
        if not self._searcher:
            self._searcher = self.index.searcher()
        return self._searcher
    
    def start_document(self):
        """Starts recording information for a new document. This should be followed by
        add_field() calls, and must be followed by an end_document() call.
        Alternatively you can use add_document() to add all fields at once.
        """
        self.segment_writer().start_document()
        
    def add_field(self, fieldname, text, stored_value = None):
        """Adds a the value of a field to the document opened with start_document().
        
        @param fieldname: The name of the field in which to index/store the text.
        @param text: The text to index.
        @type fieldname: string
        @type text: unicode
        """
        self.segment_writer().add_field(fieldname, text, stored_value = stored_value)
        
    def end_document(self):
        """
        Closes a document opened with start_document().
        """
        self.segment_writer().end_document()
    
    def add_document(self, **fields):
        """Adds all the fields of a document at once. This is an alternative to calling
        start_document(), add_field() [...], end_document().
        
        The keyword arguments map field names to the values to index/store.
        
        For fields that are both indexed and stored, you can specify an alternate
        value to store using a keyword argument in the form "_stored_<fieldname>".
        For example, if you have a field named "title" and you want to index the
        text "a b c" but store the text "e f g", use keyword arguments like this::
        
            add_document(title=u"a b c" _stored_title=u"e f g)
        """
        self.segment_writer().add_document(fields)
    
    def update_document(self, **fields):
        """Adds or replaces a document. At least one of the fields for which you
        supply values must be marked as 'unique' in the index's schema.
        
        The keyword arguments map field names to the values to index/store.
        
        For fields that are both indexed and stored, you can specify an alternate
        value to store using a keyword argument in the form "_stored_<fieldname>".
        For example, if you have a field named "title" and you want to index the
        text "a b c" but store the text "e f g", use keyword arguments like this::
        
            update_document(title=u"a b c" _stored_title=u"e f g)
        """
        
        # Check which of the supplied fields are unique
        unique_fields = [name for name, field
                         in self.index.schema.fields()
                         if name in fields and field.unique]
        if not unique_fields:
            raise IndexingError("None of the fields in %r are unique" % fields.keys())
        
        # Delete documents in which the supplied unique fields match
        searcher = self.searcher()
        for name in unique_fields:
            self.delete_by_term(name, fields[name], searcher = searcher)
        
        # Add the given fields
        self.add_document(**fields)
    
    def commit(self, mergetype = MERGE_SMALL):
        """Finishes writing and unlocks the index.
        
        @param mergetype: How to merge existing segments.
        @type mergetype: one of NO_MERGE, MERGE_SMALL, or OPTIMIZE
        """
        
        if self._segment_writer or mergetype is OPTIMIZE:
            self._merge_segments(mergetype)
        self.index.commit(self.segments)
        self._finish()
        
    def cancel(self):
        """Cancels any documents/deletions added by this object
        and unlocks the index.
        """
        self._finish()
    
    def _merge_segments(self, mergetype):
        sw = self.segment_writer()
        new_segments = mergetype(self.index, sw, self.segments)
        sw.close()
        new_segments.append(sw.segment())
        self.segments = new_segments


# Constants for "special" fields
UNIQUE_COUNT = -2
TOTAL_COUNT = -1

class SegmentWriter(object):
    """
    Do not instantiate this object directly; it is created by the IndexWriter object.
    
    Handles the actual writing of new documents to the index: writes stored fields,
    handles the posting pool, and writes out the term index.
    """
    
    class DocumentState(object):
        def __init__(self):
            self.reset()
        
        def reset(self):
            #: Whether a document is currently in progress.
            self.active = False
            #: Maps field names to stored field contents for this document
            self.stored_fields = {}
            #: Keeps track of the last field that was added.
            self.prev_fieldnum = None
    
    def __init__(self, ix, term_blocksize, doc_blocksize, vector_blocksize, name = None):
        """
        @param ix: the Index object in which to write the new segment.
        @param name: the name of the segment.
        @param blocksize: the block size to use for tables created by this writer.
        """
        
        self.index = ix
        self.schema = ix.schema
        self.storage = ix.storage
        self.name = name or ix._next_segment_name()
        
        self.max_doc = 0
        self.term_count = 0
        self.max_weight = 0
        self.doc_field_lengths = defaultdict(list)
        self.field_length_totals = defaultdict(int)
        
        # Records the state of the writer wrt start_document/end_document.
        # None == not "in" a document.
        self._doc_state = SegmentWriter.DocumentState()
        self._scorable_fields = self.schema.scorable_fields()
        
        self.pool = postpool.PostingPool()
        
        # Create a temporary segment object just so we can access
        # its *_filename attributes (so if we want to change the
        # naming convention, we only have to do it in one place).
        tempseg = index.Segment(self.name, 0, 0, 0, None)
        
        # Open files for writing
        self.term_table = self.storage.create_table(tempseg.term_filename, postings = True,
                                                    blocksize = term_blocksize)
        self.doclength_table = self.storage.create_table(tempseg.doclen_filename,
                                                         blocksize = term_blocksize)
        self.docs_table = self.storage.create_table(tempseg.docs_filename,
                                                    blocksize = doc_blocksize, compressed = 9)
        
        self.vector_table = None
        if self.schema.has_vectored_fields():
            self.vector_table = self.storage.create_table(tempseg.vector_filename,
                                                          postings = True,
                                                          stringids = True,
                                                          blocksize = vector_blocksize)
            
    def segment(self):
        """Returns an index.Segment object for the segment being written."""
        return index.Segment(self.name, self.max_doc,
                             self.term_count, self.max_weight,
                             dict(self.field_length_totals))
    
    def close(self):
        """Finishes writing the segment (flushes the posting pool out to disk) and
        closes all open files.
        """
        
        if self._doc_state.active:
            raise IndexingError("Called SegmentWriter.close() with a document still opened")
        
        self._flush_pool()
        
        self.doclength_table.add_row(UNIQUE_COUNT, self.doc_field_lengths[UNIQUE_COUNT])
        self.doclength_table.add_row(TOTAL_COUNT, self.doc_field_lengths[TOTAL_COUNT])
        for fieldnum in self._scorable_fields:
            arr = array("i", self.doc_field_lengths[fieldnum])
            self.doclength_table.add_row((fieldnum), arr)
        self.doclength_table.close()
        
{'fraction_non_alphanumeric': 0.08338294241377942, 'fraction_numerical': 0.0027781085049807517, 'mean_word_length': 2.094056974459725, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 11, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 28, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4854247', 'n_tokens_mistral': 6584, 'n_tokens_neox': 6425, 'n_words': 2318}