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