335 lines
11 KiB
Python
335 lines
11 KiB
Python
# Copyright 2012 Matt Chaput. All rights reserved.
|
|
#
|
|
# Redistribution and use in source and binary forms, with or without
|
|
# modification, are permitted provided that the following conditions are met:
|
|
#
|
|
# 1. Redistributions of source code must retain the above copyright notice,
|
|
# this list of conditions and the following disclaimer.
|
|
#
|
|
# 2. Redistributions in binary form must reproduce the above copyright
|
|
# notice, this list of conditions and the following disclaimer in the
|
|
# documentation and/or other materials provided with the distribution.
|
|
#
|
|
# THIS SOFTWARE IS PROVIDED BY MATT CHAPUT ``AS IS'' AND ANY EXPRESS OR
|
|
# IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
|
|
# MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO
|
|
# EVENT SHALL MATT CHAPUT OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
|
|
# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
|
|
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA,
|
|
# OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
|
|
# LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
|
|
# NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE,
|
|
# EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
|
#
|
|
# The views and conclusions contained in the software and documentation are
|
|
# those of the authors and should not be interpreted as representing official
|
|
# policies, either expressed or implied, of Matt Chaput.
|
|
|
|
from __future__ import with_statement
|
|
from bisect import bisect_left
|
|
from threading import Lock, RLock
|
|
|
|
from whoosh.compat import xrange
|
|
from whoosh.codec import base
|
|
from whoosh.matching import ListMatcher
|
|
from whoosh.reading import SegmentReader, TermInfo, TermNotFound
|
|
from whoosh.writing import SegmentWriter
|
|
|
|
|
|
class MemWriter(SegmentWriter):
|
|
def commit(self):
|
|
self._finalize_segment()
|
|
|
|
|
|
class MemoryCodec(base.Codec):
|
|
def __init__(self):
|
|
from whoosh.filedb.filestore import RamStorage
|
|
|
|
self.storage = RamStorage()
|
|
self.segment = MemSegment(self, "blah")
|
|
|
|
def writer(self, schema):
|
|
ix = self.storage.create_index(schema)
|
|
return MemWriter(ix, _lk=False, codec=self,
|
|
docbase=self.segment._doccount)
|
|
|
|
def reader(self, schema):
|
|
return SegmentReader(self.storage, schema, self.segment, codec=self)
|
|
|
|
def per_document_writer(self, storage, segment):
|
|
return MemPerDocWriter(self.storage, self.segment)
|
|
|
|
def field_writer(self, storage, segment):
|
|
return MemFieldWriter(self.storage, self.segment)
|
|
|
|
def per_document_reader(self, storage, segment):
|
|
return MemPerDocReader(self.storage, self.segment)
|
|
|
|
def terms_reader(self, storage, segment):
|
|
return MemTermsReader(self.storage, self.segment)
|
|
|
|
def new_segment(self, storage, indexname):
|
|
return self.segment
|
|
|
|
|
|
class MemPerDocWriter(base.PerDocWriterWithColumns):
|
|
def __init__(self, storage, segment):
|
|
self._storage = storage
|
|
self._segment = segment
|
|
self.is_closed = False
|
|
self._colwriters = {}
|
|
self._doccount = 0
|
|
|
|
def _has_column(self, fieldname):
|
|
return fieldname in self._colwriters
|
|
|
|
def _create_column(self, fieldname, column):
|
|
colfile = self._storage.create_file("%s.c" % fieldname)
|
|
self._colwriters[fieldname] = (colfile, column.writer(colfile))
|
|
|
|
def _get_column(self, fieldname):
|
|
return self._colwriters[fieldname][1]
|
|
|
|
def start_doc(self, docnum):
|
|
self._doccount += 1
|
|
self._docnum = docnum
|
|
self._stored = {}
|
|
self._lengths = {}
|
|
self._vectors = {}
|
|
|
|
def add_field(self, fieldname, fieldobj, value, length):
|
|
if value is not None:
|
|
self._stored[fieldname] = value
|
|
if length is not None:
|
|
self._lengths[fieldname] = length
|
|
|
|
def add_vector_items(self, fieldname, fieldobj, items):
|
|
self._vectors[fieldname] = tuple(items)
|
|
|
|
def finish_doc(self):
|
|
with self._segment._lock:
|
|
docnum = self._docnum
|
|
self._segment._stored[docnum] = self._stored
|
|
self._segment._lengths[docnum] = self._lengths
|
|
self._segment._vectors[docnum] = self._vectors
|
|
|
|
def close(self):
|
|
colwriters = self._colwriters
|
|
for fieldname in colwriters:
|
|
colfile, colwriter = colwriters[fieldname]
|
|
colwriter.finish(self._doccount)
|
|
colfile.close()
|
|
self.is_closed = True
|
|
|
|
|
|
class MemPerDocReader(base.PerDocumentReader):
|
|
def __init__(self, storage, segment):
|
|
self._storage = storage
|
|
self._segment = segment
|
|
|
|
def doc_count(self):
|
|
return self._segment.doc_count()
|
|
|
|
def doc_count_all(self):
|
|
return self._segment.doc_count_all()
|
|
|
|
def has_deletions(self):
|
|
return self._segment.has_deletions()
|
|
|
|
def is_deleted(self, docnum):
|
|
return self._segment.is_deleted(docnum)
|
|
|
|
def deleted_docs(self):
|
|
return self._segment.deleted_docs()
|
|
|
|
def supports_columns(self):
|
|
return True
|
|
|
|
def has_column(self, fieldname):
|
|
filename = "%s.c" % fieldname
|
|
return self._storage.file_exists(filename)
|
|
|
|
def column_reader(self, fieldname, column):
|
|
filename = "%s.c" % fieldname
|
|
colfile = self._storage.open_file(filename)
|
|
length = self._storage.file_length(filename)
|
|
return column.reader(colfile, 0, length, self._segment.doc_count_all())
|
|
|
|
def doc_field_length(self, docnum, fieldname, default=0):
|
|
return self._segment._lengths[docnum].get(fieldname, default)
|
|
|
|
def field_length(self, fieldname):
|
|
return sum(lens.get(fieldname, 0) for lens
|
|
in self._segment._lengths.values())
|
|
|
|
def min_field_length(self, fieldname):
|
|
return min(lens[fieldname] for lens in self._segment._lengths.values()
|
|
if fieldname in lens)
|
|
|
|
def max_field_length(self, fieldname):
|
|
return max(lens[fieldname] for lens in self._segment._lengths.values()
|
|
if fieldname in lens)
|
|
|
|
def has_vector(self, docnum, fieldname):
|
|
return (docnum in self._segment._vectors
|
|
and fieldname in self._segment._vectors[docnum])
|
|
|
|
def vector(self, docnum, fieldname, format_):
|
|
items = self._segment._vectors[docnum][fieldname]
|
|
ids, weights, values = zip(*items)
|
|
return ListMatcher(ids, weights, values, format_)
|
|
|
|
def stored_fields(self, docnum):
|
|
return self._segment._stored[docnum]
|
|
|
|
def close(self):
|
|
pass
|
|
|
|
|
|
class MemFieldWriter(base.FieldWriter):
|
|
def __init__(self, storage, segment):
|
|
self._storage = storage
|
|
self._segment = segment
|
|
self._fieldname = None
|
|
self._btext = None
|
|
self.is_closed = False
|
|
|
|
def start_field(self, fieldname, fieldobj):
|
|
if self._fieldname is not None:
|
|
raise Exception("Called start_field in a field")
|
|
|
|
with self._segment._lock:
|
|
invindex = self._segment._invindex
|
|
if fieldname not in invindex:
|
|
invindex[fieldname] = {}
|
|
|
|
self._fieldname = fieldname
|
|
self._fieldobj = fieldobj
|
|
|
|
def start_term(self, btext):
|
|
if self._btext is not None:
|
|
raise Exception("Called start_term in a term")
|
|
fieldname = self._fieldname
|
|
|
|
fielddict = self._segment._invindex[fieldname]
|
|
terminfos = self._segment._terminfos
|
|
with self._segment._lock:
|
|
if btext not in fielddict:
|
|
fielddict[btext] = []
|
|
|
|
if (fieldname, btext) not in terminfos:
|
|
terminfos[fieldname, btext] = TermInfo()
|
|
|
|
self._postings = fielddict[btext]
|
|
self._terminfo = terminfos[fieldname, btext]
|
|
self._btext = btext
|
|
|
|
def add(self, docnum, weight, vbytes, length):
|
|
self._postings.append((docnum, weight, vbytes))
|
|
self._terminfo.add_posting(docnum, weight, length)
|
|
|
|
def finish_term(self):
|
|
if self._btext is None:
|
|
raise Exception("Called finish_term outside a term")
|
|
|
|
self._postings = None
|
|
self._btext = None
|
|
self._terminfo = None
|
|
|
|
def finish_field(self):
|
|
if self._fieldname is None:
|
|
raise Exception("Called finish_field outside a field")
|
|
self._fieldname = None
|
|
self._fieldobj = None
|
|
|
|
def close(self):
|
|
self.is_closed = True
|
|
|
|
|
|
class MemTermsReader(base.TermsReader):
|
|
def __init__(self, storage, segment):
|
|
self._storage = storage
|
|
self._segment = segment
|
|
self._invindex = segment._invindex
|
|
|
|
def __contains__(self, term):
|
|
return term in self._segment._terminfos
|
|
|
|
def terms(self):
|
|
for fieldname in self._invindex:
|
|
for btext in self._invindex[fieldname]:
|
|
yield (fieldname, btext)
|
|
|
|
def terms_from(self, fieldname, prefix):
|
|
if fieldname not in self._invindex:
|
|
raise TermNotFound("Unknown field %r" % (fieldname,))
|
|
terms = sorted(self._invindex[fieldname])
|
|
if not terms:
|
|
return
|
|
start = bisect_left(terms, prefix)
|
|
for i in xrange(start, len(terms)):
|
|
yield (fieldname, terms[i])
|
|
|
|
def term_info(self, fieldname, text):
|
|
return self._segment._terminfos[fieldname, text]
|
|
|
|
def matcher(self, fieldname, btext, format_, scorer=None):
|
|
items = self._invindex[fieldname][btext]
|
|
ids, weights, values = zip(*items)
|
|
return ListMatcher(ids, weights, values, format_, scorer=scorer)
|
|
|
|
def indexed_field_names(self):
|
|
return self._invindex.keys()
|
|
|
|
def close(self):
|
|
pass
|
|
|
|
|
|
class MemSegment(base.Segment):
|
|
def __init__(self, codec, indexname):
|
|
base.Segment.__init__(self, indexname)
|
|
self._codec = codec
|
|
self._doccount = 0
|
|
self._stored = {}
|
|
self._lengths = {}
|
|
self._vectors = {}
|
|
self._invindex = {}
|
|
self._terminfos = {}
|
|
self._lock = Lock()
|
|
|
|
def codec(self):
|
|
return self._codec
|
|
|
|
def set_doc_count(self, doccount):
|
|
self._doccount = doccount
|
|
|
|
def doc_count(self):
|
|
return len(self._stored)
|
|
|
|
def doc_count_all(self):
|
|
return self._doccount
|
|
|
|
def delete_document(self, docnum, delete=True):
|
|
if not delete:
|
|
raise Exception("MemoryCodec can't undelete")
|
|
with self._lock:
|
|
del self._stored[docnum]
|
|
del self._lengths[docnum]
|
|
del self._vectors[docnum]
|
|
|
|
def has_deletions(self):
|
|
with self._lock:
|
|
return self._doccount - len(self._stored)
|
|
|
|
def is_deleted(self, docnum):
|
|
return docnum not in self._stored
|
|
|
|
def deleted_docs(self):
|
|
stored = self._stored
|
|
for docnum in xrange(self.doc_count_all()):
|
|
if docnum not in stored:
|
|
yield docnum
|
|
|
|
def should_assemble(self):
|
|
return False
|