lucene-java-3.5.0/lucene/contrib/benchmark/src/java/org/apache/lucene/benchmark/byTask/feeds/DocMaker.java

   1 package org.apache.lucene.benchmark.byTask.feeds;
   2
   3 /**
   4  * Licensed to the Apache Software Foundation (ASF) under one or more
   5  * contributor license agreements.  See the NOTICE file distributed with
   6  * this work for additional information regarding copyright ownership.
   7  * The ASF licenses this file to You under the Apache License, Version 2.0
   8  * (the "License"); you may not use this file except in compliance with
   9  * the License.  You may obtain a copy of the License at
  10  *
  11  *     http://www.apache.org/licenses/LICENSE-2.0
  12  *
  13  * Unless required by applicable law or agreed to in writing, software
  14  * distributed under the License is distributed on an "AS IS" BASIS,
  15  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  16  * See the License for the specific language governing permissions and
  17  * limitations under the License.
  18  */
  19
  20 import java.io.Closeable;
  21 import java.io.IOException;
  22 import java.io.UnsupportedEncodingException;
  23 import java.util.HashMap;
  24 import java.util.Calendar;
  25 import java.util.Map;
  26 import java.util.Properties;
  27 import java.util.Locale;
  28 import java.util.Random;
  29 import java.util.Date;
  30 import java.util.concurrent.atomic.AtomicInteger;
  31 import java.text.SimpleDateFormat;
  32 import java.text.ParsePosition;
  33
  34 import org.apache.lucene.benchmark.byTask.utils.Config;
  35 import org.apache.lucene.document.Document;
  36 import org.apache.lucene.document.Field;
  37 import org.apache.lucene.document.NumericField;
  38 import org.apache.lucene.document.Field.Index;
  39 import org.apache.lucene.document.Field.Store;
  40 import org.apache.lucene.document.Field.TermVector;
  41
  42 /**
  43  * Creates {@link Document} objects. Uses a {@link ContentSource} to generate
  44  * {@link DocData} objects. Supports the following parameters:
  45  * <ul>
  46  * <li><b>content.source</b> - specifies the {@link ContentSource} class to use
  47  * (default <b>SingleDocSource</b>).
  48  * <li><b>doc.stored</b> - specifies whether fields should be stored (default
  49  * <b>false</b>).
  50  * <li><b>doc.body.stored</b> - specifies whether the body field should be stored (default
  51  * = <b>doc.stored</b>).
  52  * <li><b>doc.tokenized</b> - specifies whether fields should be tokenized
  53  * (default <b>true</b>).
  54  * <li><b>doc.body.tokenized</b> - specifies whether the
  55  * body field should be tokenized (default = <b>doc.tokenized</b>).
  56  * <li><b>doc.tokenized.norms</b> - specifies whether norms should be stored in
  57  * the index or not. (default <b>false</b>).
  58  * <li><b>doc.body.tokenized.norms</b> - specifies whether norms should be
  59  * stored in the index for the body field. This can be set to true, while
  60  * <code>doc.tokenized.norms</code> is set to false, to allow norms storing just
  61  * for the body field. (default <b>true</b>).
  62  * <li><b>doc.term.vector</b> - specifies whether term vectors should be stored
  63  * for fields (default <b>false</b>).
  64  * <li><b>doc.term.vector.positions</b> - specifies whether term vectors should
  65  * be stored with positions (default <b>false</b>).
  66  * <li><b>doc.term.vector.offsets</b> - specifies whether term vectors should be
  67  * stored with offsets (default <b>false</b>).
  68  * <li><b>doc.store.body.bytes</b> - specifies whether to store the raw bytes of
  69  * the document's content in the document (default <b>false</b>).
  70  * <li><b>doc.reuse.fields</b> - specifies whether Field and Document objects
  71  * should be reused (default <b>true</b>).
  72  * <li><b>doc.index.props</b> - specifies whether the properties returned by
  73  * <li><b>doc.random.id.limit</b> - if specified, docs will be assigned random
  74  * IDs from 0 to this limit.  This is useful with UpdateDoc
  75  * for testing performance of IndexWriter.updateDocument.
  76  * {@link DocData#getProps()} will be indexed. (default <b>false</b>).
  77  * </ul>
  78  */
  79 public class DocMaker implements Closeable {
  80
  81   private static class LeftOver {
  82     public LeftOver() {}
  83     DocData docdata;
  84     int cnt;
  85   }
  86
  87   private Random r;
  88   private int updateDocIDLimit;
  89
  90   static class DocState {
  91
  92     private final Map<String,Field> fields;
  93     private final Map<String,NumericField> numericFields;
  94     private final boolean reuseFields;
  95     final Document doc;
  96     DocData docData = new DocData();
  97
  98     public DocState(boolean reuseFields, Store store, Store bodyStore, Index index, Index bodyIndex, TermVector termVector) {
  99
 100       this.reuseFields = reuseFields;
 101
 102       if (reuseFields) {
 103         fields =  new HashMap<String,Field>();
 104         numericFields = new HashMap<String,NumericField>();
 105
 106         // Initialize the map with the default fields.
 107         fields.put(BODY_FIELD, new Field(BODY_FIELD, "", bodyStore, bodyIndex, termVector));
 108         fields.put(TITLE_FIELD, new Field(TITLE_FIELD, "", store, index, termVector));
 109         fields.put(DATE_FIELD, new Field(DATE_FIELD, "", store, index, termVector));
 110         fields.put(ID_FIELD, new Field(ID_FIELD, "", Field.Store.YES, Field.Index.NOT_ANALYZED_NO_NORMS));
 111         fields.put(NAME_FIELD, new Field(NAME_FIELD, "", store, index, termVector));
 112
 113         numericFields.put(DATE_MSEC_FIELD, new NumericField(DATE_MSEC_FIELD));
 114         numericFields.put(TIME_SEC_FIELD, new NumericField(TIME_SEC_FIELD));
 115
 116         doc = new Document();
 117       } else {
 118         numericFields = null;
 119         fields = null;
 120         doc = null;
 121       }
 122     }
 123
 124     /**
 125      * Returns a field corresponding to the field name. If
 126      * <code>reuseFields</code> was set to true, then it attempts to reuse a
 127      * Field instance. If such a field does not exist, it creates a new one.
 128      */
 129     Field getField(String name, Store store, Index index, TermVector termVector) {
 130       if (!reuseFields) {
 131         return new Field(name, "", store, index, termVector);
 132       }
 133
 134       Field f = fields.get(name);
 135       if (f == null) {
 136         f = new Field(name, "", store, index, termVector);
 137         fields.put(name, f);
 138       }
 139       return f;
 140     }
 141
 142     NumericField getNumericField(String name) {
 143       if (!reuseFields) {
 144         return new NumericField(name);
 145       }
 146
 147       NumericField f = numericFields.get(name);
 148       if (f == null) {
 149         f = new NumericField(name);
 150         numericFields.put(name, f);
 151       }
 152       return f;
 153     }
 154   }
 155
 156   private boolean storeBytes = false;
 157
 158   private static class DateUtil {
 159     public SimpleDateFormat parser = new SimpleDateFormat("dd-MMM-yyyy HH:mm:ss", Locale.US);
 160     public Calendar cal = Calendar.getInstance();
 161     public ParsePosition pos = new ParsePosition(0);
 162     public DateUtil() {
 163       parser.setLenient(true);
 164     }
 165   }
 166
 167   // leftovers are thread local, because it is unsafe to share residues between threads
 168   private ThreadLocal<LeftOver> leftovr = new ThreadLocal<LeftOver>();
 169   private ThreadLocal<DocState> docState = new ThreadLocal<DocState>();
 170   private ThreadLocal<DateUtil> dateParsers = new ThreadLocal<DateUtil>();
 171
 172   public static final String BODY_FIELD = "body";
 173   public static final String TITLE_FIELD = "doctitle";
 174   public static final String DATE_FIELD = "docdate";
 175   public static final String DATE_MSEC_FIELD = "docdatenum";
 176   public static final String TIME_SEC_FIELD = "doctimesecnum";
 177   public static final String ID_FIELD = "docid";
 178   public static final String BYTES_FIELD = "bytes";
 179   public static final String NAME_FIELD = "docname";
 180
 181   protected Config config;
 182
 183   protected Store storeVal = Store.NO;
 184   protected Store bodyStoreVal = Store.NO;
 185   protected Index indexVal = Index.ANALYZED_NO_NORMS;
 186   protected Index bodyIndexVal = Index.ANALYZED;
 187   protected TermVector termVecVal = TermVector.NO;
 188
 189   protected ContentSource source;
 190   protected boolean reuseFields;
 191   protected boolean indexProperties;
 192
 193   private final AtomicInteger numDocsCreated = new AtomicInteger();
 194
 195   // create a doc
 196   // use only part of the body, modify it to keep the rest (or use all if size==0).
 197   // reset the docdata properties so they are not added more than once.
 198   private Document createDocument(DocData docData, int size, int cnt) throws UnsupportedEncodingException {
 199
 200     final DocState ds = getDocState();
 201     final Document doc = reuseFields ? ds.doc : new Document();
 202     doc.getFields().clear();
 203
 204     // Set ID_FIELD
 205     Field idField = ds.getField(ID_FIELD, storeVal, Index.NOT_ANALYZED_NO_NORMS, termVecVal);
 206     int id;
 207     if (r != null) {
 208       id = r.nextInt(updateDocIDLimit);
 209     } else {
 210       id = docData.getID();
 211       if (id == -1) {
 212         id = numDocsCreated.getAndIncrement();
 213       }
 214     }
 215     idField.setValue(Integer.toString(id));
 216     doc.add(idField);
 217
 218     // Set NAME_FIELD
 219     String name = docData.getName();
 220     if (name == null) name = "";
 221     name = cnt < 0 ? name : name + "_" + cnt;
 222     Field nameField = ds.getField(NAME_FIELD, storeVal, indexVal, termVecVal);
 223     nameField.setValue(name);
 224     doc.add(nameField);
 225
 226     // Set DATE_FIELD
 227     DateUtil util = dateParsers.get();
 228     if (util == null) {
 229       util = new DateUtil();
 230       dateParsers.set(util);
 231     }
 232     Date date = null;
 233     String dateString = docData.getDate();
 234     if (dateString != null) {
 235       util.pos.setIndex(0);
 236       date = util.parser.parse(dateString, util.pos);
 237       //System.out.println(dateString + " parsed to " + date);
 238     } else {
 239       dateString = "";
 240     }
 241     Field dateStringField = ds.getField(DATE_FIELD, storeVal, indexVal, termVecVal);
 242     dateStringField.setValue(dateString);
 243     doc.add(dateStringField);
 244
 245     if (date == null) {
 246       // just set to right now
 247       date = new Date();
 248     }
 249
 250     NumericField dateField = ds.getNumericField(DATE_MSEC_FIELD);
 251     dateField.setLongValue(date.getTime());
 252     doc.add(dateField);
 253
 254     util.cal.setTime(date);
 255     final int sec = util.cal.get(Calendar.HOUR_OF_DAY)*3600 + util.cal.get(Calendar.MINUTE)*60 + util.cal.get(Calendar.SECOND);
 256
 257     NumericField timeSecField = ds.getNumericField(TIME_SEC_FIELD);
 258     timeSecField.setIntValue(sec);
 259     doc.add(timeSecField);
 260
 261     // Set TITLE_FIELD
 262     String title = docData.getTitle();
 263     Field titleField = ds.getField(TITLE_FIELD, storeVal, indexVal, termVecVal);
 264     titleField.setValue(title == null ? "" : title);
 265     doc.add(titleField);
 266
 267     String body = docData.getBody();
 268     if (body != null && body.length() > 0) {
 269       String bdy;
 270       if (size <= 0 || size >= body.length()) {
 271         bdy = body; // use all
 272         docData.setBody(""); // nothing left
 273       } else {
 274         // attempt not to break words - if whitespace found within next 20 chars...
 275         for (int n = size - 1; n < size + 20 && n < body.length(); n++) {
 276           if (Character.isWhitespace(body.charAt(n))) {
 277             size = n;
 278             break;
 279           }
 280         }
 281         bdy = body.substring(0, size); // use part
 282         docData.setBody(body.substring(size)); // some left
 283       }
 284       Field bodyField = ds.getField(BODY_FIELD, bodyStoreVal, bodyIndexVal, termVecVal);
 285       bodyField.setValue(bdy);
 286       doc.add(bodyField);
 287
 288       if (storeBytes) {
 289         Field bytesField = ds.getField(BYTES_FIELD, Store.YES, Index.NOT_ANALYZED_NO_NORMS, TermVector.NO);
 290         bytesField.setValue(bdy.getBytes("UTF-8"));
 291         doc.add(bytesField);
 292       }
 293     }
 294
 295     if (indexProperties) {
 296       Properties props = docData.getProps();
 297       if (props != null) {
 298         for (final Map.Entry<Object,Object> entry : props.entrySet()) {
 299           Field f = ds.getField((String) entry.getKey(), storeVal, indexVal, termVecVal);
 300           f.setValue((String) entry.getValue());
 301           doc.add(f);
 302         }
 303         docData.setProps(null);
 304       }
 305     }
 306
 307     //System.out.println("============== Created doc "+numDocsCreated+" :\n"+doc+"\n==========");
 308     return doc;
 309   }
 310
 311   private void resetLeftovers() {
 312     leftovr.set(null);
 313   }
 314
 315   protected DocState getDocState() {
 316     DocState ds = docState.get();
 317     if (ds == null) {
 318       ds = new DocState(reuseFields, storeVal, bodyStoreVal, indexVal, bodyIndexVal, termVecVal);
 319       docState.set(ds);
 320     }
 321     return ds;
 322   }
 323
 324   /**
 325    * Closes the {@link DocMaker}. The base implementation closes the
 326    * {@link ContentSource}, and it can be overridden to do more work (but make
 327    * sure to call super.close()).
 328    */
 329   public void close() throws IOException {
 330     source.close();
 331   }
 332
 333   /**
 334    * Returns the number of bytes generated by the content source since last
 335    * reset.
 336    */
 337   public synchronized long getBytesCount() {
 338     return source.getBytesCount();
 339   }
 340
 341   /**
 342    * Returns the total number of bytes that were generated by the content source
 343    * defined to that doc maker.
 344    */
 345   public long getTotalBytesCount() {
 346     return source.getTotalBytesCount();
 347   }
 348
 349   /**
 350    * Creates a {@link Document} object ready for indexing. This method uses the
 351    * {@link ContentSource} to get the next document from the source, and creates
 352    * a {@link Document} object from the returned fields. If
 353    * <code>reuseFields</code> was set to true, it will reuse {@link Document}
 354    * and {@link Field} instances.
 355    */
 356   public Document makeDocument() throws Exception {
 357     resetLeftovers();
 358     DocData docData = source.getNextDocData(getDocState().docData);
 359     Document doc = createDocument(docData, 0, -1);
 360     return doc;
 361   }
 362
 363   /**
 364    * Same as {@link #makeDocument()}, only this method creates a document of the
 365    * given size input by <code>size</code>.
 366    */
 367   public Document makeDocument(int size) throws Exception {
 368     LeftOver lvr = leftovr.get();
 369     if (lvr == null || lvr.docdata == null || lvr.docdata.getBody() == null
 370         || lvr.docdata.getBody().length() == 0) {
 371       resetLeftovers();
 372     }
 373     DocData docData = getDocState().docData;
 374     DocData dd = (lvr == null ? source.getNextDocData(docData) : lvr.docdata);
 375     int cnt = (lvr == null ? 0 : lvr.cnt);
 376     while (dd.getBody() == null || dd.getBody().length() < size) {
 377       DocData dd2 = dd;
 378       dd = source.getNextDocData(new DocData());
 379       cnt = 0;
 380       dd.setBody(dd2.getBody() + dd.getBody());
 381     }
 382     Document doc = createDocument(dd, size, cnt);
 383     if (dd.getBody() == null || dd.getBody().length() == 0) {
 384       resetLeftovers();
 385     } else {
 386       if (lvr == null) {
 387         lvr = new LeftOver();
 388         leftovr.set(lvr);
 389       }
 390       lvr.docdata = dd;
 391       lvr.cnt = ++cnt;
 392     }
 393     return doc;
 394   }
 395
 396   /** Reset inputs so that the test run would behave, input wise, as if it just started. */
 397   public synchronized void resetInputs() throws IOException {
 398     source.printStatistics("docs");
 399     // re-initiate since properties by round may have changed.
 400     setConfig(config);
 401     source.resetInputs();
 402     numDocsCreated.set(0);
 403     resetLeftovers();
 404   }
 405
 406   /** Set the configuration parameters of this doc maker. */
 407   public void setConfig(Config config) {
 408     this.config = config;
 409     try {
 410       String sourceClass = config.get("content.source", "org.apache.lucene.benchmark.byTask.feeds.SingleDocSource");
 411       source = Class.forName(sourceClass).asSubclass(ContentSource.class).newInstance();
 412       source.setConfig(config);
 413     } catch (Exception e) {
 414       // Should not get here. Throw runtime exception.
 415       throw new RuntimeException(e);
 416     }
 417
 418     boolean stored = config.get("doc.stored", false);
 419     boolean bodyStored = config.get("doc.body.stored", stored);
 420     boolean tokenized = config.get("doc.tokenized", true);
 421     boolean bodyTokenized = config.get("doc.body.tokenized", tokenized);
 422     boolean norms = config.get("doc.tokenized.norms", false);
 423     boolean bodyNorms = config.get("doc.body.tokenized.norms", true);
 424     boolean termVec = config.get("doc.term.vector", false);
 425     storeVal = (stored ? Field.Store.YES : Field.Store.NO);
 426     bodyStoreVal = (bodyStored ? Field.Store.YES : Field.Store.NO);
 427     if (tokenized) {
 428       indexVal = norms ? Index.ANALYZED : Index.ANALYZED_NO_NORMS;
 429     } else {
 430       indexVal = norms ? Index.NOT_ANALYZED : Index.NOT_ANALYZED_NO_NORMS;
 431     }
 432
 433     if (bodyTokenized) {
 434       bodyIndexVal = bodyNorms ? Index.ANALYZED : Index.ANALYZED_NO_NORMS;
 435     } else {
 436       bodyIndexVal = bodyNorms ? Index.NOT_ANALYZED : Index.NOT_ANALYZED_NO_NORMS;
 437     }
 438
 439     boolean termVecPositions = config.get("doc.term.vector.positions", false);
 440     boolean termVecOffsets = config.get("doc.term.vector.offsets", false);
 441     if (termVecPositions && termVecOffsets) {
 442       termVecVal = TermVector.WITH_POSITIONS_OFFSETS;
 443     } else if (termVecPositions) {
 444       termVecVal = TermVector.WITH_POSITIONS;
 445     } else if (termVecOffsets) {
 446       termVecVal = TermVector.WITH_OFFSETS;
 447     } else if (termVec) {
 448       termVecVal = TermVector.YES;
 449     } else {
 450       termVecVal = TermVector.NO;
 451     }
 452     storeBytes = config.get("doc.store.body.bytes", false);
 453
 454     reuseFields = config.get("doc.reuse.fields", true);
 455
 456     // In a multi-rounds run, it is important to reset DocState since settings
 457     // of fields may change between rounds, and this is the only way to reset
 458     // the cache of all threads.
 459     docState = new ThreadLocal<DocState>();
 460
 461     indexProperties = config.get("doc.index.props", false);
 462
 463     updateDocIDLimit = config.get("doc.random.id.limit", -1);
 464     if (updateDocIDLimit != -1) {
 465       r = new Random(179);
 466     }
 467   }
 468
 469 }