lucene/benchmark/src/java/org/apache/lucene/benchmark/byTask/feeds/DirContentSource.java - lucene-solr - Git at Google

 /*
  * Licensed to the Apache Software Foundation (ASF) under one or more
  * contributor license agreements.  See the NOTICE file distributed with
  * this work for additional information regarding copyright ownership.
  * The ASF licenses this file to You under the Apache License, Version 2.0
  * (the "License"); you may not use this file except in compliance with
  * the License.  You may obtain a copy of the License at
  *
  *     http://www.apache.org/licenses/LICENSE-2.0
  *
  * Unless required by applicable law or agreed to in writing, software
  * distributed under the License is distributed on an "AS IS" BASIS,
  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
 package org.apache.lucene.benchmark.byTask.feeds;


 import org.apache.lucene.benchmark.byTask.utils.Config;

 import java.io.BufferedReader;
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
 import java.nio.file.DirectoryStream;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.text.DateFormat;
 import java.text.ParsePosition;
 import java.text.SimpleDateFormat;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Date;
 import java.util.List;
 import java.util.Locale;
 import java.util.Stack;

 /**
  * A {@link ContentSource} using the Dir collection for its input. Supports
  * the following configuration parameters (on top of {@link ContentSource}):
  * <ul>
  * <li><b>work.dir</b> - specifies the working directory. Required if "docs.dir"
  * denotes a relative path (<b>default=work</b>).
  * <li><b>docs.dir</b> - specifies the directory the Dir collection. Can be set
  * to a relative path if "work.dir" is also specified (<b>default=dir-out</b>).
  * </ul>
  */
 public class DirContentSource extends ContentSource {

   private static final class DateFormatInfo {
     DateFormat df;
     ParsePosition pos;
   }

   /**
    * Iterator over the files in the directory
    */
   public static class Iterator implements java.util.Iterator<Path> {

     static class Comparator implements java.util.Comparator<Path> {
       @Override
       public int compare(Path _a, Path _b) {
         String a = _a.toString();
         String b = _b.toString();
         int diff = a.length() - b.length();

         if (diff > 0) {
           while (diff-- > 0) {
             b = "0" + b;
           }
         } else if (diff < 0) {
           diff = -diff;
           while (diff-- > 0) {
             a = "0" + a;
           }
         }

         /* note it's reversed because we're going to push,
            which reverses again */
         return b.compareTo(a);
       }
     }

     int count = 0;

     Stack<Path> stack = new Stack<>();

     /* this seems silly ... there must be a better way ...
        not that this is good, but can it matter? */

     Comparator c = new Comparator();

     public Iterator(Path f) throws IOException {
       push(f);
     }

     void find() throws IOException {
       if (stack.empty()) {
         return;
       }
       if (!Files.isDirectory(stack.peek())) {
         return;
       }
       Path f = stack.pop();
       push(f);
     }

     void push(Path f) throws IOException {
       try (DirectoryStream<Path> stream = Files.newDirectoryStream(f)) {
         List<Path> found = new ArrayList<>();
         for (Path p : stream) {
           if (Files.isDirectory(p)) {
             found.add(p);
           }
         }
         push(found.toArray(new Path[found.size()]));
       }
       try (DirectoryStream<Path> stream = Files.newDirectoryStream(f, "*.txt")) {
         List<Path> found = new ArrayList<>();
         for (Path p : stream) {
           found.add(p);
         }
         push(found.toArray(new Path[found.size()]));
       }
       find();
     }

     void push(Path[] files) {
       Arrays.sort(files, c);
       for(int i = 0; i < files.length; i++) {
         // System.err.println("push " + files[i]);
         stack.push(files[i]);
       }
     }

     public int getCount(){
       return count;
     }

     @Override
     public boolean hasNext() {
       return stack.size() > 0;
     }

     @Override
     public Path next() {
       assert hasNext();
       count++;
       Path object = stack.pop();
       // System.err.println("pop " + object);
       try {
         find();
       } catch (IOException e) {
         throw new RuntimeException(e);
       }
       return object;
     }

     @Override
     public void remove() {
       throw new RuntimeException("cannot");
     }

   }

   private ThreadLocal<DateFormatInfo> dateFormat = new ThreadLocal<>();
   private Path dataDir = null;
   private int iteration = 0;
   private Iterator inputFiles = null;

   // get/initiate a thread-local simple date format (must do so
   // because SimpleDateFormat is not thread-safe).
   private DateFormatInfo getDateFormatInfo() {
     DateFormatInfo dfi = dateFormat.get();
     if (dfi == null) {
       dfi = new DateFormatInfo();
       dfi.pos = new ParsePosition(0);
       // date format: 30-MAR-1987 14:22:36.87
       dfi.df = new SimpleDateFormat("dd-MMM-yyyy kk:mm:ss.SSS", Locale.ENGLISH);
       dfi.df.setLenient(true);
       dateFormat.set(dfi);
     }
     return dfi;
   }

   private Date parseDate(String dateStr) {
     DateFormatInfo dfi = getDateFormatInfo();
     dfi.pos.setIndex(0);
     dfi.pos.setErrorIndex(-1);
     return dfi.df.parse(dateStr.trim(), dfi.pos);
   }

   @Override
   public void close() throws IOException {
     inputFiles = null;
   }

   @Override
   public DocData getNextDocData(DocData docData) throws NoMoreDataException, IOException {
     Path f = null;
     String name = null;
     synchronized (this) {
       if (!inputFiles.hasNext()) {
         // exhausted files, start a new round, unless forever set to false.
         if (!forever) {
           throw new NoMoreDataException();
         }
         inputFiles = new Iterator(dataDir);
         iteration++;
       }
       f = inputFiles.next();
       // System.err.println(f);
       name = f.toRealPath()+"_"+iteration;
     }

     BufferedReader reader = Files.newBufferedReader(f, StandardCharsets.UTF_8);
     String line = null;
     //First line is the date, 3rd is the title, rest is body
     String dateStr = reader.readLine();
     reader.readLine();//skip an empty line
     String title = reader.readLine();
     reader.readLine();//skip an empty line
     StringBuilder bodyBuf = new StringBuilder(1024);
     while ((line = reader.readLine()) != null) {
       bodyBuf.append(line).append(' ');
     }
     reader.close();
     addBytes(Files.size(f));

     Date date = parseDate(dateStr);

     docData.clear();
     docData.setName(name);
     docData.setBody(bodyBuf.toString());
     docData.setTitle(title);
     docData.setDate(date);
     return docData;
   }

   @Override
   public synchronized void resetInputs() throws IOException {
     super.resetInputs();
     inputFiles = new Iterator(dataDir);
     iteration = 0;
   }

   @Override
   public void setConfig(Config config) {
     super.setConfig(config);

     Path workDir = Paths.get(config.get("work.dir", "work"));
     String d = config.get("docs.dir", "dir-out");
     dataDir = Paths.get(d);
     if (!dataDir.isAbsolute()) {
       dataDir = workDir.resolve(d);
     }

     try {
       inputFiles = new Iterator(dataDir);
     } catch (IOException e) {
       throw new RuntimeException(e);
     }

     if (inputFiles == null) {
       throw new RuntimeException("No txt files in dataDir: " + dataDir.toAbsolutePath());
     }
   }

 }
	/*
	* Licensed to the Apache Software Foundation (ASF) under one or more
	* contributor license agreements. See the NOTICE file distributed with
	* this work for additional information regarding copyright ownership.
	* The ASF licenses this file to You under the Apache License, Version 2.0
	* (the "License"); you may not use this file except in compliance with
	* the License. You may obtain a copy of the License at
	*
	* http://www.apache.org/licenses/LICENSE-2.0
	*
	* Unless required by applicable law or agreed to in writing, software
	* distributed under the License is distributed on an "AS IS" BASIS,
	* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	* See the License for the specific language governing permissions and
	* limitations under the License.
	*/
	package org.apache.lucene.benchmark.byTask.feeds;


	import org.apache.lucene.benchmark.byTask.utils.Config;

	import java.io.BufferedReader;
	import java.io.IOException;
	import java.nio.charset.StandardCharsets;
	import java.nio.file.DirectoryStream;
	import java.nio.file.Files;
	import java.nio.file.Path;
	import java.nio.file.Paths;
	import java.text.DateFormat;
	import java.text.ParsePosition;
	import java.text.SimpleDateFormat;
	import java.util.ArrayList;
	import java.util.Arrays;
	import java.util.Date;
	import java.util.List;
	import java.util.Locale;
	import java.util.Stack;

	/**
	* A {@link ContentSource} using the Dir collection for its input. Supports
	* the following configuration parameters (on top of {@link ContentSource}):
	* <ul>
	* <li><b>work.dir</b> - specifies the working directory. Required if "docs.dir"
	* denotes a relative path (<b>default=work</b>).
	* <li><b>docs.dir</b> - specifies the directory the Dir collection. Can be set
	* to a relative path if "work.dir" is also specified (<b>default=dir-out</b>).
	* </ul>
	*/
	public class DirContentSource extends ContentSource {

	private static final class DateFormatInfo {
	DateFormat df;
	ParsePosition pos;
	}

	/**
	* Iterator over the files in the directory
	*/
	public static class Iterator implements java.util.Iterator<Path> {

	static class Comparator implements java.util.Comparator<Path> {
	@Override
	public int compare(Path _a, Path _b) {
	String a = _a.toString();
	String b = _b.toString();
	int diff = a.length() - b.length();

	if (diff > 0) {
	while (diff-- > 0) {
	b = "0" + b;
	}
	} else if (diff < 0) {
	diff = -diff;
	while (diff-- > 0) {
	a = "0" + a;
	}
	}

	/* note it's reversed because we're going to push,
	which reverses again */
	return b.compareTo(a);
	}
	}

	int count = 0;

	Stack<Path> stack = new Stack<>();

	/* this seems silly ... there must be a better way ...
	not that this is good, but can it matter? */

	Comparator c = new Comparator();

	public Iterator(Path f) throws IOException {
	push(f);
	}

	void find() throws IOException {
	if (stack.empty()) {
	return;
	}
	if (!Files.isDirectory(stack.peek())) {
	return;
	}
	Path f = stack.pop();
	push(f);
	}

	void push(Path f) throws IOException {
	try (DirectoryStream<Path> stream = Files.newDirectoryStream(f)) {
	List<Path> found = new ArrayList<>();
	for (Path p : stream) {
	if (Files.isDirectory(p)) {
	found.add(p);
	}
	}
	push(found.toArray(new Path[found.size()]));
	}
	try (DirectoryStream<Path> stream = Files.newDirectoryStream(f, "*.txt")) {
	List<Path> found = new ArrayList<>();
	for (Path p : stream) {
	found.add(p);
	}
	push(found.toArray(new Path[found.size()]));
	}
	find();
	}

	void push(Path[] files) {
	Arrays.sort(files, c);
	for(int i = 0; i < files.length; i++) {
	// System.err.println("push " + files[i]);
	stack.push(files[i]);
	}
	}

	public int getCount(){
	return count;
	}

	@Override
	public boolean hasNext() {
	return stack.size() > 0;
	}

	@Override
	public Path next() {
	assert hasNext();
	count++;
	Path object = stack.pop();
	// System.err.println("pop " + object);
	try {
	find();
	} catch (IOException e) {
	throw new RuntimeException(e);
	}
	return object;
	}

	@Override
	public void remove() {
	throw new RuntimeException("cannot");
	}

	}

	private ThreadLocal<DateFormatInfo> dateFormat = new ThreadLocal<>();
	private Path dataDir = null;
	private int iteration = 0;
	private Iterator inputFiles = null;

	// get/initiate a thread-local simple date format (must do so
	// because SimpleDateFormat is not thread-safe).
	private DateFormatInfo getDateFormatInfo() {
	DateFormatInfo dfi = dateFormat.get();
	if (dfi == null) {
	dfi = new DateFormatInfo();
	dfi.pos = new ParsePosition(0);
	// date format: 30-MAR-1987 14:22:36.87
	dfi.df = new SimpleDateFormat("dd-MMM-yyyy kk:mm:ss.SSS", Locale.ENGLISH);
	dfi.df.setLenient(true);
	dateFormat.set(dfi);
	}
	return dfi;
	}

	private Date parseDate(String dateStr) {
	DateFormatInfo dfi = getDateFormatInfo();
	dfi.pos.setIndex(0);
	dfi.pos.setErrorIndex(-1);
	return dfi.df.parse(dateStr.trim(), dfi.pos);
	}

	@Override
	public void close() throws IOException {
	inputFiles = null;
	}

	@Override
	public DocData getNextDocData(DocData docData) throws NoMoreDataException, IOException {
	Path f = null;
	String name = null;
	synchronized (this) {
	if (!inputFiles.hasNext()) {
	// exhausted files, start a new round, unless forever set to false.
	if (!forever) {
	throw new NoMoreDataException();
	}
	inputFiles = new Iterator(dataDir);
	iteration++;
	}
	f = inputFiles.next();
	// System.err.println(f);
	name = f.toRealPath()+"_"+iteration;
	}

	BufferedReader reader = Files.newBufferedReader(f, StandardCharsets.UTF_8);
	String line = null;
	//First line is the date, 3rd is the title, rest is body
	String dateStr = reader.readLine();
	reader.readLine();//skip an empty line
	String title = reader.readLine();
	reader.readLine();//skip an empty line
	StringBuilder bodyBuf = new StringBuilder(1024);
	while ((line = reader.readLine()) != null) {
	bodyBuf.append(line).append(' ');
	}
	reader.close();
	addBytes(Files.size(f));

	Date date = parseDate(dateStr);

	docData.clear();
	docData.setName(name);
	docData.setBody(bodyBuf.toString());
	docData.setTitle(title);
	docData.setDate(date);
	return docData;
	}

	@Override
	public synchronized void resetInputs() throws IOException {
	super.resetInputs();
	inputFiles = new Iterator(dataDir);
	iteration = 0;
	}

	@Override
	public void setConfig(Config config) {
	super.setConfig(config);

	Path workDir = Paths.get(config.get("work.dir", "work"));
	String d = config.get("docs.dir", "dir-out");
	dataDir = Paths.get(d);
	if (!dataDir.isAbsolute()) {
	dataDir = workDir.resolve(d);
	}

	try {
	inputFiles = new Iterator(dataDir);
	} catch (IOException e) {
	throw new RuntimeException(e);
	}

	if (inputFiles == null) {
	throw new RuntimeException("No txt files in dataDir: " + dataDir.toAbsolutePath());
	}
	}

	}