apps/search/index.py

   1 # -*- coding: utf-8 -*-
   2 from django.conf import settings
   3 from lucene import SimpleFSDirectory, IndexWriter, File, Field, \
   4     NumericField, Version, Document, JavaError, IndexSearcher, \
   5     QueryParser, Term, PerFieldAnalyzerWrapper, \
   6     SimpleAnalyzer, PolishAnalyzer, ArrayList, \
   7     KeywordAnalyzer, NumericRangeQuery, BooleanQuery, \
   8     BlockJoinQuery, BlockJoinCollector, TermsFilter, \
   9     HashSet, BooleanClause, Term, CharTermAttribute, \
  10     PhraseQuery, StringReader, TermQuery, BlockJoinQuery, \
  11     Sort, Integer
  12     # KeywordAnalyzer
  13 import sys
  14 import os
  15 import errno
  16 from librarian import dcparser
  17 from librarian.parser import WLDocument
  18 import catalogue.models
  19 from multiprocessing.pool import ThreadPool
  20 from threading import current_thread
  21 import atexit
  22
  23
  24 class WLAnalyzer(PerFieldAnalyzerWrapper):
  25     def __init__(self):
  26         polish = PolishAnalyzer(Version.LUCENE_34)
  27         simple = SimpleAnalyzer(Version.LUCENE_34)
  28         keyword = KeywordAnalyzer(Version.LUCENE_34)
  29         # not sure if needed: there's NOT_ANALYZED meaning basically the same
  30
  31         PerFieldAnalyzerWrapper.__init__(self, polish)
  32
  33         self.addAnalyzer("tags", simple)
  34         self.addAnalyzer("technical_editors", simple)
  35         self.addAnalyzer("editors", simple)
  36         self.addAnalyzer("url", keyword)
  37         self.addAnalyzer("source_url", keyword)
  38         self.addAnalyzer("source_name", simple)
  39         self.addAnalyzer("publisher", simple)
  40         self.addAnalyzer("author", simple)
  41         self.addAnalyzer("is_book", keyword)
  42
  43         #self.addanalyzer("fragment_anchor", keyword)
  44
  45
  46 class IndexStore(object):
  47     def __init__(self):
  48         self.make_index_dir()
  49         self.store = SimpleFSDirectory(File(settings.SEARCH_INDEX))
  50
  51     def make_index_dir(self):
  52         try:
  53             os.makedirs(settings.SEARCH_INDEX)
  54         except OSError as exc:
  55             if exc.errno == errno.EEXIST:
  56                 pass
  57             else: raise
  58
  59
  60 class Index(IndexStore):
  61     def __init__(self, analyzer=None):
  62         IndexStore.__init__(self)
  63         self.index = None
  64         if not analyzer:
  65             analyzer = WLAnalyzer()
  66         self.analyzer = analyzer
  67
  68     def open(self, analyzer=None):
  69         if self.index:
  70             raise Exception("Index is already opened")
  71         self.index = IndexWriter(self.store, self.analyzer,\
  72                                  IndexWriter.MaxFieldLength.LIMITED)
  73         return self.index
  74
  75     def close(self):
  76         self.index.optimize()
  77         self.index.close()
  78         self.index = None
  79
  80     def remove_book(self, book):
  81         q = NumericRangeQuery.newIntRange("book_id", book.id, book.id, True,True)
  82         self.index.deleteDocuments(q)
  83
  84     def index_book(self, book, overwrite=True):
  85         if overwrite:
  86             self.remove_book(book)
  87
  88
  89         doc = self.extract_metadata(book)
  90         parts = self.extract_content(book)
  91         block = ArrayList().of_(Document)
  92
  93         for p in parts:
  94             block.add(p)
  95         block.add(doc)
  96         self.index.addDocuments(block)
  97
  98     master_tags = [
  99         'opowiadanie',
 100         'powiesc',
 101         'dramat_wierszowany_l',
 102         'dramat_wierszowany_lp',
 103         'dramat_wspolczesny', 'liryka_l', 'liryka_lp',
 104         'wywiad'
 105         ]
 106
 107     skip_header_tags = ['autor_utworu', 'nazwa_utworu', 'dzielo_nadrzedne']
 108
 109     def create_book_doc(self, book):
 110         """
 111         Create a lucene document connected to the book
 112         """
 113         doc = Document()
 114         doc.add(NumericField("book_id", Field.Store.YES, True).setIntValue(book.id))
 115         if book.parent is not None:
 116             doc.add(NumericField("parent_id", Field.Store.YES, True).setIntValue(book.parent.id))
 117         return doc
 118
 119     def extract_metadata(self, book):
 120         book_info = dcparser.parse(book.xml_file)
 121
 122         print("extract metadata for book %s id=%d, thread%d" % (book.slug, book.id, current_thread().ident))
 123
 124         doc = self.create_book_doc(book)
 125         doc.add(Field("slug", book.slug, Field.Store.NO, Field.Index.ANALYZED_NO_NORMS))
 126         doc.add(Field("tags", ','.join([t.name for t in book.tags]), Field.Store.NO, Field.Index.ANALYZED))
 127         doc.add(Field("is_book", 'true', Field.Store.NO, Field.Index.NOT_ANALYZED))
 128
 129         # validator, name
 130         for field in dcparser.BookInfo.FIELDS:
 131             if hasattr(book_info, field.name):
 132                 if not getattr(book_info, field.name):
 133                     continue
 134                 # since no type information is available, we use validator
 135                 type_indicator = field.validator
 136                 if type_indicator == dcparser.as_unicode:
 137                     s = getattr(book_info, field.name)
 138                     if field.multiple:
 139                         s = ', '.join(s)
 140                     try:
 141                         doc.add(Field(field.name, s, Field.Store.NO, Field.Index.ANALYZED))
 142                     except JavaError as je:
 143                         raise Exception("failed to add field: %s = '%s', %s(%s)" % (field.name, s, je.message, je.args))
 144                 elif type_indicator == dcparser.as_person:
 145                     p = getattr(book_info, field.name)
 146                     if isinstance(p, dcparser.Person):
 147                         persons = unicode(p)
 148                     else:
 149                         persons = ', '.join(map(unicode, p))
 150                     doc.add(Field(field.name, persons, Field.Store.NO, Field.Index.ANALYZED))
 151                 elif type_indicator == dcparser.as_date:
 152                     dt = getattr(book_info, field.name)
 153                     doc.add(Field(field.name, "%04d%02d%02d" % (dt.year, dt.month, dt.day), Field.Store.NO, Field.Index.NOT_ANALYZED))
 154         return doc
 155
 156     def get_master(self, root):
 157         for master in root.iter():
 158             if master.tag in self.master_tags:
 159                 return master
 160
 161
 162     def extract_content(self, book):
 163         wld = WLDocument.from_file(book.xml_file.path)
 164         root = wld.edoc.getroot()
 165
 166         # first we build a sequence of top-level items.
 167         # book_id
 168         # header_index - the 0-indexed position of header element.
 169         # content
 170         master = self.get_master(root)
 171         if master is None:
 172             return []
 173
 174         header_docs = []
 175         for header, position in zip(list(master), range(len(master))):
 176             if header.tag in self.skip_header_tags:
 177                 continue
 178             doc = self.create_book_doc(book)
 179             doc.add(NumericField("header_index", Field.Store.YES, True).setIntValue(position))
 180             doc.add(Field("header_type", header.tag, Field.Store.YES, Field.Index.NOT_ANALYZED))
 181             content = u' '.join([t for t in header.itertext()])
 182             doc.add(Field("content", content, Field.Store.YES, Field.Index.ANALYZED))
 183             header_docs.append(doc)
 184
 185         def walker(node):
 186             yield node, None
 187             for child in list(node):
 188                 for b, e in walker(child):
 189                     yield b, e
 190             yield None, node
 191             return
 192
 193         # Then we create a document for each fragments
 194         # fragment_anchor - the anchor
 195         # themes - list of themes [not indexed]
 196         fragment_docs = []
 197         # will contain (framgent id -> { content: [], themes: [] }
 198         fragments = {}
 199         for start, end in walker(master):
 200             if start is not None and start.tag == 'begin':
 201                 fid = start.attrib['id'][1:]
 202                 fragments[fid] = {'content': [], 'themes': []}
 203                 fragments[fid]['content'].append(start.tail)
 204             elif start is not None and start.tag == 'motyw':
 205                 fid = start.attrib['id'][1:]
 206                 fragments[fid]['themes'].append(start.text)
 207                 fragments[fid]['content'].append(start.tail)
 208             elif start is not None and start.tag == 'end':
 209                 fid = start.attrib['id'][1:]
 210                 if fid not in fragments:
 211                     continue  # a broken <end> node, skip it
 212                 frag = fragments[fid]
 213                 del fragments[fid]
 214
 215                 def jstr(l):
 216                     return u' '.join(map(
 217                         lambda x: x == None and u'(none)' or unicode(x),
 218                         l))
 219
 220                 doc = self.create_book_doc(book)
 221                 doc.add(Field("fragment_anchor", fid,
 222                               Field.Store.YES, Field.Index.NOT_ANALYZED))
 223                 doc.add(Field("content",
 224                               u' '.join(filter(lambda s: s is not None, frag['content'])),
 225                               Field.Store.YES, Field.Index.ANALYZED))
 226                 doc.add(Field("themes",
 227                               u' '.join(filter(lambda s: s is not None, frag['themes'])),
 228                               Field.Store.NO, Field.Index.ANALYZED))
 229
 230                 fragment_docs.append(doc)
 231             elif start is not None:
 232                 for frag in fragments.values():
 233                     frag['content'].append(start.text)
 234             elif end is not None:
 235                 for frag in fragments.values():
 236                     frag['content'].append(end.tail)
 237
 238         return header_docs + fragment_docs
 239
 240     def __enter__(self):
 241         self.open()
 242         return self
 243
 244     def __exit__(self, type, value, tb):
 245         self.close()
 246
 247
 248 class ReusableIndex(Index):
 249     """
 250     Works like index, but does not close/optimize Lucene index
 251     until program exit (uses atexit hook).
 252     This is usefull for importbooks command.
 253
 254     if you cannot rely on atexit, use ReusableIndex.close_reusable() yourself.
 255     """
 256     index = None
 257     pool = None
 258     pool_jobs = None
 259
 260     def open(self, analyzer=None, threads=4):
 261         if ReusableIndex.index is not None:
 262             self.index = ReusableIndex.index
 263         else:
 264             print("opening index")
 265             ReusableIndex.pool = ThreadPool(threads)
 266             ReusableIndex.pool_jobs = []
 267             Index.open(self, analyzer)
 268             ReusableIndex.index = self.index
 269             atexit.register(ReusableIndex.close_reusable)
 270
 271     def index_book(self, *args, **kw):
 272         job = ReusableIndex.pool.apply_async(Index.index_book, (self,)+ args, kw)
 273         ReusableIndex.pool_jobs.append(job)
 274
 275     @staticmethod
 276     def close_reusable():
 277         if ReusableIndex.index is not None:
 278             print("closing index")
 279             for job in ReusableIndex.pool_jobs:
 280                 job.wait()
 281             ReusableIndex.pool.close()
 282
 283             ReusableIndex.index.optimize()
 284             ReusableIndex.index.close()
 285             ReusableIndex.index = None
 286
 287     def close(self):
 288         pass
 289
 290
 291 class Search(IndexStore):
 292     def __init__(self, default_field="content"):
 293         IndexStore.__init__(self)
 294         self.analyzer = PolishAnalyzer(Version.LUCENE_34)
 295         ## self.analyzer = WLAnalyzer()
 296         self.searcher = IndexSearcher(self.store, True)
 297         self.parser = QueryParser(Version.LUCENE_34, default_field,
 298                                   self.analyzer)
 299
 300         self.parent_filter = TermsFilter()
 301         self.parent_filter.addTerm(Term("is_book", "true"))
 302
 303     def query(self, query):
 304         return self.parser.parse(query)
 305
 306     def wrapjoins(self, query, fields=[]):
 307         """
 308         This functions modifies the query in a recursive way,
 309         so Term and Phrase Queries contained, which match
 310         provided fields are wrapped in a BlockJoinQuery,
 311         and so delegated to children documents.
 312         """
 313         if BooleanQuery.instance_(query):
 314             qs = BooleanQuery.cast_(query)
 315             for clause in qs:
 316                 clause = BooleanClause.cast_(clause)
 317                 clause.setQuery(self.wrapjoins(clause.getQuery(), fields))
 318             return qs
 319         else:
 320             termset = HashSet()
 321             query.extractTerms(termset)
 322             for t in termset:
 323                 t = Term.cast_(t)
 324                 if t.field() not in fields:
 325                     return query
 326             return BlockJoinQuery(query, self.parent_filter,
 327                                   BlockJoinQuery.ScoreMode.Total)
 328
 329     def simple_search(self, query, max_results=50):
 330         """Returns (books, total_hits)
 331         """
 332
 333         tops = self.searcher.search(self.query(query), max_results)
 334         bks = []
 335         for found in tops.scoreDocs:
 336             doc = self.searcher.doc(found.doc)
 337             bks.append(catalogue.models.Book.objects.get(id=doc.get("book_id")))
 338         return (bks, tops.totalHits)
 339
 340     def search(self, query, max_results=50):
 341         query = self.query(query)
 342         query = self.wrapjoins(query, ["content", "themes"])
 343
 344         tops = self.searcher.search(query, max_results)
 345         bks = []
 346         for found in tops.scoreDocs:
 347             doc = self.searcher.doc(found.doc)
 348             bks.append(catalogue.models.Book.objects.get(id=doc.get("book_id")))
 349         return (bks, tops.totalHits)
 350
 351     def bsearch(self, query, max_results=50):
 352         q = self.query(query)
 353         bjq = BlockJoinQuery(q, self.parent_filter, BlockJoinQuery.ScoreMode.Avg)
 354
 355         tops = self.searcher.search(bjq, max_results)
 356         bks = []
 357         for found in tops.scoreDocs:
 358             doc = self.searcher.doc(found.doc)
 359             bks.append(catalogue.models.Book.objects.get(id=doc.get("book_id")))
 360         return (bks, tops.totalHits)
 361
 362 # TokenStream tokenStream = analyzer.tokenStream(fieldName, reader);
 363 # OffsetAttribute offsetAttribute = tokenStream.getAttribute(OffsetAttribute.class);
 364 # CharTermAttribute charTermAttribute = tokenStream.getAttribute(CharTermAttribute.class);
 365
 366 # while (tokenStream.incrementToken()) {
 367 #     int startOffset = offsetAttribute.startOffset();
 368 #     int endOffset = offsetAttribute.endOffset();
 369 #     String term = charTermAttribute.toString();
 370 # }
 371
 372
 373 class MultiSearch(Search):
 374     """Class capable of IMDb-like searching"""
 375     def get_tokens(self, queryreader):
 376         if isinstance(queryreader, str) or isinstance(queryreader, unicode):
 377             queryreader = StringReader(queryreader)
 378         queryreader.reset()
 379         tokens = self.analyzer.reusableTokenStream('content', queryreader)
 380         toks = []
 381         while tokens.incrementToken():
 382             cta = tokens.getAttribute(CharTermAttribute.class_)
 383             toks.append(cta.toString())
 384         return toks
 385
 386     def make_phrase(self, tokens, field='content', slop=2):
 387         phrase = PhraseQuery()
 388         phrase.setSlop(slop)
 389         for t in tokens:
 390             term = Term(field, t)
 391             phrase.add(term)
 392         return phrase
 393
 394     def make_term_query(self, tokens, field='content', modal=BooleanClause.Occur.SHOULD):
 395         q = BooleanQuery()
 396         for t in tokens:
 397             term = Term(field, t)
 398             q.add(BooleanClause(TermQuery(term), modal))
 399         return q
 400
 401     def content_query(self, query):
 402         return BlockJoinQuery(query, self.parent_filter,
 403                               BlockJoinQuery.ScoreMode.Total)
 404
 405     def search_perfect(self, tokens, max_results=20):
 406         qrys = [self.make_phrase(tokens, field=fld) for fld in ['author', 'title', 'content']]
 407
 408         books = []
 409         for q in qrys:
 410             top = self.searcher.search(q, max_results)
 411             for found in top.scoreDocs:
 412                 book_info = self.searcher.doc(found.doc)
 413                 books.append((found.score, catalogue.models.Book.objects.get(id=book_info.get("book_id")), []))
 414         return books
 415
 416     def search_everywhere(self, tokens, max_results=20):
 417         q = BooleanQuery()
 418         in_meta = BooleanQuery()
 419         in_content = BooleanQuery()
 420
 421         for fld in ['themes', 'content']:
 422             in_content.add(BooleanClause(self.make_term_query(tokens, field=fld), BooleanClause.Occur.SHOULD))
 423
 424         for fld in ['author', 'title', 'epochs', 'genres', 'kinds']:
 425             in_meta.add(BooleanClause(self.make_term_query(tokens, field=fld), BooleanClause.Occur.SHOULD))
 426
 427         q.add(BooleanClause(in_meta, BooleanClause.Occur.MUST))
 428         in_content_join = self.content_query(in_content)
 429         q.add(BooleanClause(in_content_join, BooleanClause.Occur.MUST))
 430
 431         collector = BlockJoinCollector(Sort.RELEVANCE, 100, True, True)
 432
 433         self.searcher.search(q, collector)
 434
 435         books = []
 436
 437         top_groups = collector.getTopGroups(in_content_join, Sort.RELEVANCE, 0, max_results, 0, True)
 438         if top_groups:
 439             for grp in top_groups.groups:
 440                 doc_id = Integer.cast_(grp.groupValue).intValue()
 441                 book_data = self.searcher.doc(doc_id)
 442                 book = catalogue.models.Book.objects.get(id=book_data.get("book_id"))
 443                 parts = []
 444                 for part in grp.scoreDocs:
 445                     part_data = self.searcher.doc(part.doc)
 446                     header_type = part_data.get("header_type")
 447                     if header_type:
 448                         parts.append((part.score, {"header": header_type, "position": int(part_data.get("header_index"))}))
 449                     fragment = part_data.get("fragment_anchor")
 450                     if fragment:
 451                         fragment = book.fragments.get(anchor=fragment)
 452                         parts.append((part.score, {"fragment": fragment}))
 453                 books.append((grp.maxScore, book, parts))
 454
 455         return books
 456
 457
 458     def multisearch(self, query, max_results=50):
 459         """
 460         Search strategy:
 461         - (phrase) OR -> content
 462                       -> title
 463                       -> author
 464         - (keywords)  -> author
 465                       -> motyw
 466                       -> tags
 467                       -> content
 468         """
 469         # queryreader = StringReader(query)
 470         # tokens = self.get_tokens(queryreader)
 471
 472         # top_level = BooleanQuery()
 473         # Should = BooleanClause.Occur.SHOULD
 474
 475         # phrase_level = BooleanQuery()
 476         # phrase_level.setBoost(1.3)
 477
 478         # p_content = self.make_phrase(tokens, joined=True)
 479         # p_title = self.make_phrase(tokens, 'title')
 480         # p_author = self.make_phrase(tokens, 'author')
 481
 482         # phrase_level.add(BooleanClause(p_content, Should))
 483         # phrase_level.add(BooleanClause(p_title, Should))
 484         # phrase_level.add(BooleanClause(p_author, Should))
 485
 486         # kw_level = BooleanQuery()
 487
 488         # kw_level.add(self.make_term_query(tokens, 'author'), Should)
 489         # j_themes = self.make_term_query(tokens, 'themes', joined=True)
 490         # kw_level.add(j_themes, Should)
 491         # kw_level.add(self.make_term_query(tokens, 'tags'), Should)
 492         # j_con = self.make_term_query(tokens, joined=True)
 493         # kw_level.add(j_con, Should)
 494
 495         # top_level.add(BooleanClause(phrase_level, Should))
 496         # top_level.add(BooleanClause(kw_level, Should))
 497
 498         return None
 499
 500
 501     def do_search(self, query, max_results=50, collector=None):
 502         tops = self.searcher.search(query, max_results)
 503         #tops = self.searcher.search(p_content, max_results)
 504
 505         bks = []
 506         for found in tops.scoreDocs:
 507             doc = self.searcher.doc(found.doc)
 508             b = catalogue.models.Book.objects.get(id=doc.get("book_id"))
 509             bks.append(b)
 510             print "%s (%d) -> %f" % (b, b.id, found.score)
 511         return (bks, tops.totalHits)