lucene-java-3.5.0/lucene/contrib/benchmark/src/java/org/apache/lucene/benchmark/byTask/feeds/DirContentSource.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 org.apache.lucene.benchmark.byTask.utils.Config;
  21
  22 import java.io.BufferedReader;
  23 import java.io.File;
  24 import java.io.FileFilter;
  25 import java.io.FileReader;
  26 import java.io.IOException;
  27 import java.text.DateFormat;
  28 import java.text.ParsePosition;
  29 import java.text.SimpleDateFormat;
  30 import java.util.Arrays;
  31 import java.util.Date;
  32 import java.util.Locale;
  33 import java.util.Stack;
  34
  35 /**
  36  * A {@link ContentSource} using the Dir collection for its input. Supports
  37  * the following configuration parameters (on top of {@link ContentSource}):
  38  * <ul>
  39  * <li><b>work.dir</b> - specifies the working directory. Required if "docs.dir"
  40  * denotes a relative path (<b>default=work</b>).
  41  * <li><b>docs.dir</b> - specifies the directory the Dir collection. Can be set
  42  * to a relative path if "work.dir" is also specified (<b>default=dir-out</b>).
  43  * </ul>
  44  */
  45 public class DirContentSource extends ContentSource {
  46
  47   private static final class DateFormatInfo {
  48     public DateFormatInfo() {}
  49     DateFormat df;
  50     ParsePosition pos;
  51   }
  52
  53   public static class Iterator implements java.util.Iterator<File> {
  54
  55     static class Comparator implements java.util.Comparator<File> {
  56       public int compare(File _a, File _b) {
  57         String a = _a.toString();
  58         String b = _b.toString();
  59         int diff = a.length() - b.length();
  60
  61         if (diff > 0) {
  62           while (diff-- > 0) {
  63             b = "0" + b;
  64           }
  65         } else if (diff < 0) {
  66           diff = -diff;
  67           while (diff-- > 0) {
  68             a = "0" + a;
  69           }
  70         }
  71
  72         /* note it's reversed because we're going to push,
  73            which reverses again */
  74         return b.compareTo(a);
  75       }
  76     }
  77
  78     int count = 0;
  79
  80     Stack<File> stack = new Stack<File>();
  81
  82     /* this seems silly ... there must be a better way ...
  83        not that this is good, but can it matter? */
  84
  85     Comparator c = new Comparator();
  86
  87     public Iterator(File f) {
  88       push(f);
  89     }
  90
  91     void find() {
  92       if (stack.empty()) {
  93         return;
  94       }
  95       if (!(stack.peek()).isDirectory()) {
  96         return;
  97       }
  98       File f = stack.pop();
  99       push(f);
 100     }
 101
 102     void push(File f) {
 103       push(f.listFiles(new FileFilter() {
 104
 105         public boolean accept(File file) {
 106           return file.isDirectory();
 107         }
 108       }));
 109       push(f.listFiles(new FileFilter() {
 110
 111         public boolean accept(File file) {
 112           return file.getName().endsWith(".txt");
 113         }
 114       }));
 115       find();
 116     }
 117
 118     void push(File[] files) {
 119       Arrays.sort(files, c);
 120       for(int i = 0; i < files.length; i++) {
 121         // System.err.println("push " + files[i]);
 122         stack.push(files[i]);
 123       }
 124     }
 125
 126     public int getCount(){
 127       return count;
 128     }
 129
 130     public boolean hasNext() {
 131       return stack.size() > 0;
 132     }
 133
 134     public File next() {
 135       assert hasNext();
 136       count++;
 137       File object = stack.pop();
 138       // System.err.println("pop " + object);
 139       find();
 140       return object;
 141     }
 142
 143     public void remove() {
 144       throw new RuntimeException("cannot");
 145     }
 146
 147   }
 148
 149   private ThreadLocal<DateFormatInfo> dateFormat = new ThreadLocal<DateFormatInfo>();
 150   private File dataDir = null;
 151   private int iteration = 0;
 152   private Iterator inputFiles = null;
 153
 154   // get/initiate a thread-local simple date format (must do so
 155   // because SimpleDateFormat is not thread-safe).
 156   private DateFormatInfo getDateFormatInfo() {
 157     DateFormatInfo dfi = dateFormat.get();
 158     if (dfi == null) {
 159       dfi = new DateFormatInfo();
 160       dfi.pos = new ParsePosition(0);
 161       // date format: 30-MAR-1987 14:22:36.87
 162       dfi.df = new SimpleDateFormat("dd-MMM-yyyy kk:mm:ss.SSS", Locale.US);
 163       dfi.df.setLenient(true);
 164       dateFormat.set(dfi);
 165     }
 166     return dfi;
 167   }
 168
 169   private Date parseDate(String dateStr) {
 170     DateFormatInfo dfi = getDateFormatInfo();
 171     dfi.pos.setIndex(0);
 172     dfi.pos.setErrorIndex(-1);
 173     return dfi.df.parse(dateStr.trim(), dfi.pos);
 174   }
 175
 176   @Override
 177   public void close() throws IOException {
 178     inputFiles = null;
 179   }
 180
 181   @Override
 182   public DocData getNextDocData(DocData docData) throws NoMoreDataException, IOException {
 183     File f = null;
 184     String name = null;
 185     synchronized (this) {
 186       if (!inputFiles.hasNext()) {
 187         // exhausted files, start a new round, unless forever set to false.
 188         if (!forever) {
 189           throw new NoMoreDataException();
 190         }
 191         inputFiles = new Iterator(dataDir);
 192         iteration++;
 193       }
 194       f = inputFiles.next();
 195       // System.err.println(f);
 196       name = f.getCanonicalPath()+"_"+iteration;
 197     }
 198
 199     BufferedReader reader = new BufferedReader(new FileReader(f));
 200     String line = null;
 201     //First line is the date, 3rd is the title, rest is body
 202     String dateStr = reader.readLine();
 203     reader.readLine();//skip an empty line
 204     String title = reader.readLine();
 205     reader.readLine();//skip an empty line
 206     StringBuilder bodyBuf = new StringBuilder(1024);
 207     while ((line = reader.readLine()) != null) {
 208       bodyBuf.append(line).append(' ');
 209     }
 210     reader.close();
 211     addBytes(f.length());
 212
 213     Date date = parseDate(dateStr);
 214
 215     docData.clear();
 216     docData.setName(name);
 217     docData.setBody(bodyBuf.toString());
 218     docData.setTitle(title);
 219     docData.setDate(date);
 220     return docData;
 221   }
 222
 223   @Override
 224   public synchronized void resetInputs() throws IOException {
 225     super.resetInputs();
 226     inputFiles = new Iterator(dataDir);
 227     iteration = 0;
 228   }
 229
 230   @Override
 231   public void setConfig(Config config) {
 232     super.setConfig(config);
 233
 234     File workDir = new File(config.get("work.dir", "work"));
 235     String d = config.get("docs.dir", "dir-out");
 236     dataDir = new File(d);
 237     if (!dataDir.isAbsolute()) {
 238       dataDir = new File(workDir, d);
 239     }
 240
 241     inputFiles = new Iterator(dataDir);
 242
 243     if (inputFiles == null) {
 244       throw new RuntimeException("No txt files in dataDir: " + dataDir.getAbsolutePath());
 245     }
 246   }
 247
 248 }