org.fbk.cit.hlt.thewikimachine.index.PageTextIndexer Maven / Gradle / Ivy

Go to download
/*
 * Copyright (2013) Fondazione Bruno Kessler (http://www.fbk.eu/)
 *
 * Licensed 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.fbk.cit.hlt.thewikimachine.index;

import org.apache.commons.cli.*;
import org.apache.log4j.Logger;
import org.apache.log4j.PropertyConfigurator;
import org.apache.lucene.document.Document;
import org.apache.lucene.document.Field;
import org.fbk.cit.hlt.thewikimachine.index.util.AbstractIndexer;
import org.fbk.cit.hlt.thewikimachine.index.util.SingleValueIndexer;
import org.fbk.cit.hlt.thewikimachine.util.StringTable;
import org.xerial.snappy.SnappyInputStream;

import java.io.*;
import java.text.DecimalFormat;
import java.util.Date;
import java.util.regex.Pattern;


/**
 * Created with IntelliJ IDEA.
 * User: giuliano
 * Date: 1/24/13
 * Time: 6:47 PM
 * To change this template use File | Settings | File Templates.
 */
public class PageTextIndexer extends AbstractIndexer {
	/**
	 * Define a static logger variable so that it references the
	 * Logger instance named PageTextIndexer.
	 */
	static Logger logger = Logger.getLogger(PageTextIndexer.class.getName());

	protected static Pattern spacePattern = Pattern.compile(StringTable.SPACE);

	public static final String PAGE_FIELD_NAME = "page";

	public static final String TEXT_FIELD_NAME = "text";

	public static final int PAGE_COLUMN_INDEX = 0;

	public static final boolean DEFAULT_FILE_COMPRESS = false;

	protected static DecimalFormat df = new DecimalFormat("###,###,###,###");


	public PageTextIndexer(String indexName) throws IOException {
		super(indexName);
	}

	public void index(String fileName) throws IOException {
		index(new File(fileName), DEFAULT_FILE_COMPRESS);
	}

	public void index(String fileName, boolean compress) throws IOException {
		index(new File(fileName), compress);
	}

	public void index(File file) throws IOException {
		index(file, DEFAULT_FILE_COMPRESS);
	}

	public void index(File file, boolean compress) throws IOException {
		logger.info("indexing " + file + "...");
		long begin = System.currentTimeMillis(), end = 0;
		LineNumberReader lnr;
		if (compress) {
			lnr = new LineNumberReader(new InputStreamReader(new SnappyInputStream(new FileInputStream(file)), "UTF-8"));
		}
		else {
			lnr = new LineNumberReader(new InputStreamReader(new FileInputStream(file), "UTF-8"));
		}

		String line;
		int tot = 0, count = 0;
		String[] t;
		logger.info("tot\tcount\ttime\tdate");

		while ((line = lnr.readLine()) != null) {
			t = spacePattern.split(line);
			if (t.length > PAGE_COLUMN_INDEX) {
				add(t);
				count++;
			}
			tot++;

			//if (tot > 1000000)
			//	break;

			if ((tot % notificationPoint) == 0) {
				end = System.currentTimeMillis();
				logger.info(df.format(tot) + "\t" + df.format(count) + "\t" + df.format(end - begin) + "\t" + new Date());
				begin = System.currentTimeMillis();
			}
		}

		end = System.currentTimeMillis();
		logger.info(df.format(tot) + " lines indexed: " + df.format(count) + "\t" + df.format(end - begin) + " ms " + new Date());
		lnr.close();
	}

	public void add(String[] t) {

		if (t.length > 1) {
			Document doc = new Document();
			try {
				doc.add(new Field(PAGE_FIELD_NAME, t[PAGE_COLUMN_INDEX], Field.Store.YES, Field.Index.NOT_ANALYZED));
				doc.add(new Field(TEXT_FIELD_NAME, toByte(t), Field.Store.YES));

				indexWriter.addDocument(doc);
			} catch (IOException e) {
				logger.error(e);
			}
		}
	}

	private byte[] toByte(String[] s) throws IOException {
		ByteArrayOutputStream byteStream = new ByteArrayOutputStream(1024);
		DataOutputStream dataStream = new DataOutputStream(byteStream);
		// length
		dataStream.writeInt(s.length - 1);
		for (int i = 1; i < s.length; i++) {
			dataStream.writeUTF(s[i]);
		}
		return byteStream.toByteArray();
	}

	public static void main(String args[]) throws Exception {
		String logConfig = System.getProperty("log-config");
		if (logConfig == null) {
			logConfig = "configuration/log-config.txt";
		}


		PropertyConfigurator.configure(logConfig);
		Options options = new Options();
		try {
			Option indexNameOpt = OptionBuilder.withArgName("index").hasArg().withDescription("create an index with the specified name").isRequired().withLongOpt("index").create("i");
			Option inputFileOpt = OptionBuilder.withArgName("file").hasArg().withDescription("read the PAGE_COLUMN_INDEX/value pairs to index from the specified file").withLongOpt("file").create("f");
			//Option keyFieldNameOpt = OptionBuilder.withArgName("PAGE_COLUMN_INDEX-field-name").hasArg().withDescription("use the specified name for the field PAGE_COLUMN_INDEX").withLongOpt("PAGE_COLUMN_INDEX-field-name").create("k");
			//Option valueFieldNameOpt = OptionBuilder.withArgName("value-field-name").hasArg().withDescription("use the specified name for the field value").withLongOpt("value-field-name").create("v");
			//Option freqFileOpt = OptionBuilder.withArgName("PAGE_COLUMN_INDEX-freq").withDescription("PAGE_COLUMN_INDEX frequency file").withLongOpt("PAGE_COLUMN_INDEX-freq").create("f");

			options.addOption("h", "help", false, "print this message");
			options.addOption("v", "version", false, "output version information and exit");

			options.addOption(indexNameOpt);
			options.addOption(inputFileOpt);
			//options.addOption(keyFieldNameOpt);
			//options.addOption(valueFieldNameOpt);
			CommandLineParser parser = new PosixParser();
			CommandLine line = parser.parse(options, args);

			if (line.hasOption("help") || line.hasOption("version")) {
				throw new ParseException("");
			}

			PageTextIndexer pageTextIndexer = new PageTextIndexer(line.getOptionValue("index"));
			pageTextIndexer.index(line.getOptionValue("file"));
			pageTextIndexer.close();
		} catch (ParseException e) {
			// oops, something went wrong
			if (e.getMessage().length() > 0) {
				System.out.println("Parsing failed: " + e.getMessage() + "\n");
			}
			HelpFormatter formatter = new HelpFormatter();
			formatter.printHelp(400, "java -cp dist/thewikimachine.jar org.fbk.cit.hlt.thewikimachine.index.PageTextIndexer", "\n", options, "\n", true);
		}
	}
}