1 package org.apache.lucene.benchmark.byTask.tasks;
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
11 * http://www.apache.org/licenses/LICENSE-2.0
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.
20 import java.io.BufferedWriter;
22 import java.io.OutputStream;
23 import java.io.OutputStreamWriter;
24 import java.io.PrintWriter;
25 import java.util.Arrays;
26 import java.util.HashSet;
27 import java.util.regex.Matcher;
28 import java.util.regex.Pattern;
30 import org.apache.lucene.benchmark.byTask.PerfRunData;
31 import org.apache.lucene.benchmark.byTask.feeds.DocMaker;
32 import org.apache.lucene.benchmark.byTask.utils.Config;
33 import org.apache.lucene.benchmark.byTask.utils.StreamUtils;
34 import org.apache.lucene.document.Document;
35 import org.apache.lucene.document.Fieldable;
38 * A task which writes documents, one line per document. Each line is in the
39 * following format: title <TAB> date <TAB> body. The output of this
40 * task can be consumed by
41 * {@link org.apache.lucene.benchmark.byTask.feeds.LineDocSource} and is intended
42 * to save the IO overhead of opening a file per document to be indexed.
44 * The format of the output is set according to the output file extension.
45 * Compression is recommended when the output file is expected to be large.
46 * See info on file extensions in
47 * {@link org.apache.lucene.benchmark.byTask.utils.StreamUtils.Type}
49 * Supports the following parameters:
51 * <li><b>line.file.out</b> - the name of the file to write the output to. That
52 * parameter is mandatory. <b>NOTE:</b> the file is re-created.
53 * <li><b>line.fields</b> - which fields should be written in each line.
54 * (optional, default: {@link #DEFAULT_FIELDS}).
55 * <li><b>sufficient.fields</b> - list of field names, separated by comma, which,
56 * if all of them are missing, the document will be skipped. For example, to require
57 * that at least one of f1,f2 is not empty, specify: "f1,f2" in this field. To specify
58 * that no field is required, i.e. that even empty docs should be emitted, specify <b>","</b>.
59 * (optional, default: {@link #DEFAULT_SUFFICIENT_FIELDS}).
61 * <b>NOTE:</b> this class is not thread-safe and if used by multiple threads the
62 * output is unspecified (as all will write to the same output file in a
63 * non-synchronized way).
65 public class WriteLineDocTask extends PerfTask {
67 public static final String FIELDS_HEADER_INDICATOR = "FIELDS_HEADER_INDICATOR###";
69 public final static char SEP = '\t';
72 * Fields to be written by default
74 public static final String[] DEFAULT_FIELDS = new String[] {
81 * Default fields which at least one of them is required to not skip the doc.
83 public static final String DEFAULT_SUFFICIENT_FIELDS = DocMaker.TITLE_FIELD +',' + DocMaker.BODY_FIELD;
85 private int docSize = 0;
86 private PrintWriter lineFileOut = null;
87 private DocMaker docMaker;
88 private ThreadLocal<StringBuilder> threadBuffer = new ThreadLocal<StringBuilder>();
89 private ThreadLocal<Matcher> threadNormalizer = new ThreadLocal<Matcher>();
90 private final String[] fieldsToWrite;
91 private final boolean[] sufficientFields;
92 private final boolean checkSufficientFields;
94 public WriteLineDocTask(PerfRunData runData) throws Exception {
96 Config config = runData.getConfig();
97 String fname = config.get("line.file.out", null);
99 throw new IllegalArgumentException("line.file.out must be set");
101 OutputStream out = StreamUtils.outputStream(new File(fname));
102 lineFileOut = new PrintWriter(new BufferedWriter(new OutputStreamWriter(out, "UTF-8"), StreamUtils.BUFFER_SIZE));
103 docMaker = runData.getDocMaker();
106 String f2r = config.get("line.fields",null);
108 fieldsToWrite = DEFAULT_FIELDS;
110 if (f2r.indexOf(SEP)>=0) {
111 throw new IllegalArgumentException("line.fields "+f2r+" should not contain the separator char: "+SEP);
113 fieldsToWrite = f2r.split(",");
116 // init sufficient fields
117 sufficientFields = new boolean[fieldsToWrite.length];
118 String suff = config.get("sufficient.fields",DEFAULT_SUFFICIENT_FIELDS);
119 if (",".equals(suff)) {
120 checkSufficientFields = false;
122 checkSufficientFields = true;
123 HashSet<String> sf = new HashSet<String>(Arrays.asList(suff.split(",")));
124 for (int i=0; i<fieldsToWrite.length; i++) {
125 if (sf.contains(fieldsToWrite[i])) {
126 sufficientFields[i] = true;
135 * Write a header to the lines file - indicating how to read the file later
137 private void writeHeader() {
138 StringBuilder sb = threadBuffer.get();
140 sb = new StringBuilder();
141 threadBuffer.set(sb);
144 sb.append(FIELDS_HEADER_INDICATOR);
145 for (String f : fieldsToWrite) {
146 sb.append(SEP).append(f);
148 lineFileOut.println(sb.toString());
152 protected String getLogMessage(int recsCount) {
153 return "Wrote " + recsCount + " line docs";
157 public int doLogic() throws Exception {
158 Document doc = docSize > 0 ? docMaker.makeDocument(docSize) : docMaker.makeDocument();
160 Matcher matcher = threadNormalizer.get();
161 if (matcher == null) {
162 matcher = Pattern.compile("[\t\r\n]+").matcher("");
163 threadNormalizer.set(matcher);
166 StringBuilder sb = threadBuffer.get();
168 sb = new StringBuilder();
169 threadBuffer.set(sb);
173 boolean sufficient = !checkSufficientFields;
174 for (int i=0; i<fieldsToWrite.length; i++) {
175 Fieldable f = doc.getFieldable(fieldsToWrite[i]);
176 String text = f == null ? "" : matcher.reset(f.stringValue()).replaceAll(" ").trim();
177 sb.append(text).append(SEP);
178 sufficient |= text.length()>0 && sufficientFields[i];
181 sb.setLength(sb.length()-1); // remove redundant last separator
182 // lineFileOut is a PrintWriter, which synchronizes internally in println.
183 lineFileOut.println(sb.toString());
190 public void close() throws Exception {
196 * Set the params (docSize only)
197 * @param params docSize, or 0 for no limit.
200 public void setParams(String params) {
201 super.setParams(params);
202 docSize = (int) Float.parseFloat(params);
206 public boolean supportsParams() {