001/*
002 * Licensed to the Apache Software Foundation (ASF) under one or more
003 * contributor license agreements.  See the NOTICE file distributed with
004 * this work for additional information regarding copyright ownership.
005 * The ASF licenses this file to You under the Apache License, Version 2.0
006 * (the "License"); you may not use this file except in compliance with
007 * the License.  You may obtain a copy of the License at
008 *
009 *     http://www.apache.org/licenses/LICENSE-2.0
010 *
011 * Unless required by applicable law or agreed to in writing, software
012 * distributed under the License is distributed on an "AS IS" BASIS,
013 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
014 * See the License for the specific language governing permissions and
015 * limitations under the License.
016 */
017package org.apache.lucene.demo;
018
019import java.io.BufferedReader;
020import java.io.IOException;
021import java.io.InputStreamReader;
022import java.nio.charset.StandardCharsets;
023import java.nio.file.Files;
024import java.nio.file.Paths;
025import java.util.ArrayList;
026import java.util.Date;
027import java.util.List;
028import org.apache.lucene.analysis.Analyzer;
029import org.apache.lucene.analysis.standard.StandardAnalyzer;
030import org.apache.lucene.demo.knn.DemoEmbeddings;
031import org.apache.lucene.demo.knn.KnnVectorDict;
032import org.apache.lucene.document.Document;
033import org.apache.lucene.index.DirectoryReader;
034import org.apache.lucene.index.Term;
035import org.apache.lucene.queryparser.classic.QueryParser;
036import org.apache.lucene.search.BooleanClause;
037import org.apache.lucene.search.BooleanQuery;
038import org.apache.lucene.search.IndexSearcher;
039import org.apache.lucene.search.KnnVectorQuery;
040import org.apache.lucene.search.Query;
041import org.apache.lucene.search.QueryVisitor;
042import org.apache.lucene.search.ScoreDoc;
043import org.apache.lucene.search.TopDocs;
044import org.apache.lucene.store.FSDirectory;
045
046/** Simple command-line based search demo. */
047public class SearchFiles {
048
049  private SearchFiles() {}
050
051  /** Simple command-line based search demo. */
052  public static void main(String[] args) throws Exception {
053    String usage =
054        "Usage:\tjava org.apache.lucene.demo.SearchFiles [-index dir] [-field f] [-repeat n] [-queries file] [-query string] [-raw] [-paging hitsPerPage] [-knn_vector knnHits]\n\nSee http://lucene.apache.org/core/9_0_0/demo/ for details.";
055    if (args.length > 0 && ("-h".equals(args[0]) || "-help".equals(args[0]))) {
056      System.out.println(usage);
057      System.exit(0);
058    }
059
060    String index = "index";
061    String field = "contents";
062    String queries = null;
063    int repeat = 0;
064    boolean raw = false;
065    int knnVectors = 0;
066    String queryString = null;
067    int hitsPerPage = 10;
068
069    for (int i = 0; i < args.length; i++) {
070      switch (args[i]) {
071        case "-index":
072          index = args[++i];
073          break;
074        case "-field":
075          field = args[++i];
076          break;
077        case "-queries":
078          queries = args[++i];
079          break;
080        case "-query":
081          queryString = args[++i];
082          break;
083        case "-repeat":
084          repeat = Integer.parseInt(args[++i]);
085          break;
086        case "-raw":
087          raw = true;
088          break;
089        case "-paging":
090          hitsPerPage = Integer.parseInt(args[++i]);
091          if (hitsPerPage <= 0) {
092            System.err.println("There must be at least 1 hit per page.");
093            System.exit(1);
094          }
095          break;
096        case "-knn_vector":
097          knnVectors = Integer.parseInt(args[++i]);
098          break;
099        default:
100          System.err.println("Unknown argument: " + args[i]);
101          System.exit(1);
102      }
103    }
104
105    DirectoryReader reader = DirectoryReader.open(FSDirectory.open(Paths.get(index)));
106    IndexSearcher searcher = new IndexSearcher(reader);
107    Analyzer analyzer = new StandardAnalyzer();
108    KnnVectorDict vectorDict = null;
109    if (knnVectors > 0) {
110      vectorDict = new KnnVectorDict(reader.directory(), IndexFiles.KNN_DICT);
111    }
112    BufferedReader in;
113    if (queries != null) {
114      in = Files.newBufferedReader(Paths.get(queries), StandardCharsets.UTF_8);
115    } else {
116      in = new BufferedReader(new InputStreamReader(System.in, StandardCharsets.UTF_8));
117    }
118    QueryParser parser = new QueryParser(field, analyzer);
119    while (true) {
120      if (queries == null && queryString == null) { // prompt the user
121        System.out.println("Enter query: ");
122      }
123
124      String line = queryString != null ? queryString : in.readLine();
125
126      if (line == null || line.length() == -1) {
127        break;
128      }
129
130      line = line.trim();
131      if (line.length() == 0) {
132        break;
133      }
134
135      Query query = parser.parse(line);
136      if (knnVectors > 0) {
137        query = addSemanticQuery(query, vectorDict, knnVectors);
138      }
139      System.out.println("Searching for: " + query.toString(field));
140
141      if (repeat > 0) { // repeat & time as benchmark
142        Date start = new Date();
143        for (int i = 0; i < repeat; i++) {
144          searcher.search(query, 100);
145        }
146        Date end = new Date();
147        System.out.println("Time: " + (end.getTime() - start.getTime()) + "ms");
148      }
149
150      doPagingSearch(in, searcher, query, hitsPerPage, raw, queries == null && queryString == null);
151
152      if (queryString != null) {
153        break;
154      }
155    }
156    if (vectorDict != null) {
157      vectorDict.close();
158    }
159    reader.close();
160  }
161
162  /**
163   * This demonstrates a typical paging search scenario, where the search engine presents pages of
164   * size n to the user. The user can then go to the next page if interested in the next hits.
165   *
166   * <p>When the query is executed for the first time, then only enough results are collected to
167   * fill 5 result pages. If the user wants to page beyond this limit, then the query is executed
168   * another time and all hits are collected.
169   */
170  public static void doPagingSearch(
171      BufferedReader in,
172      IndexSearcher searcher,
173      Query query,
174      int hitsPerPage,
175      boolean raw,
176      boolean interactive)
177      throws IOException {
178
179    // Collect enough docs to show 5 pages
180    TopDocs results = searcher.search(query, 5 * hitsPerPage);
181    ScoreDoc[] hits = results.scoreDocs;
182
183    int numTotalHits = Math.toIntExact(results.totalHits.value);
184    System.out.println(numTotalHits + " total matching documents");
185
186    int start = 0;
187    int end = Math.min(numTotalHits, hitsPerPage);
188
189    while (true) {
190      if (end > hits.length) {
191        System.out.println(
192            "Only results 1 - "
193                + hits.length
194                + " of "
195                + numTotalHits
196                + " total matching documents collected.");
197        System.out.println("Collect more (y/n) ?");
198        String line = in.readLine();
199        if (line.length() == 0 || line.charAt(0) == 'n') {
200          break;
201        }
202
203        hits = searcher.search(query, numTotalHits).scoreDocs;
204      }
205
206      end = Math.min(hits.length, start + hitsPerPage);
207
208      for (int i = start; i < end; i++) {
209        if (raw) { // output raw format
210          System.out.println("doc=" + hits[i].doc + " score=" + hits[i].score);
211          continue;
212        }
213
214        Document doc = searcher.doc(hits[i].doc);
215        String path = doc.get("path");
216        if (path != null) {
217          System.out.println((i + 1) + ". " + path);
218          String title = doc.get("title");
219          if (title != null) {
220            System.out.println("   Title: " + doc.get("title"));
221          }
222        } else {
223          System.out.println((i + 1) + ". " + "No path for this document");
224        }
225      }
226
227      if (!interactive || end == 0) {
228        break;
229      }
230
231      if (numTotalHits >= end) {
232        boolean quit = false;
233        while (true) {
234          System.out.print("Press ");
235          if (start - hitsPerPage >= 0) {
236            System.out.print("(p)revious page, ");
237          }
238          if (start + hitsPerPage < numTotalHits) {
239            System.out.print("(n)ext page, ");
240          }
241          System.out.println("(q)uit or enter number to jump to a page.");
242
243          String line = in.readLine();
244          if (line.length() == 0 || line.charAt(0) == 'q') {
245            quit = true;
246            break;
247          }
248          if (line.charAt(0) == 'p') {
249            start = Math.max(0, start - hitsPerPage);
250            break;
251          } else if (line.charAt(0) == 'n') {
252            if (start + hitsPerPage < numTotalHits) {
253              start += hitsPerPage;
254            }
255            break;
256          } else {
257            int page = Integer.parseInt(line);
258            if ((page - 1) * hitsPerPage < numTotalHits) {
259              start = (page - 1) * hitsPerPage;
260              break;
261            } else {
262              System.out.println("No such page");
263            }
264          }
265        }
266        if (quit) break;
267        end = Math.min(numTotalHits, start + hitsPerPage);
268      }
269    }
270  }
271
272  private static Query addSemanticQuery(Query query, KnnVectorDict vectorDict, int k)
273      throws IOException {
274    StringBuilder semanticQueryText = new StringBuilder();
275    QueryFieldTermExtractor termExtractor = new QueryFieldTermExtractor("contents");
276    query.visit(termExtractor);
277    for (String term : termExtractor.terms) {
278      semanticQueryText.append(term).append(' ');
279    }
280    if (semanticQueryText.length() > 0) {
281      KnnVectorQuery knnQuery =
282          new KnnVectorQuery(
283              "contents-vector",
284              new DemoEmbeddings(vectorDict).computeEmbedding(semanticQueryText.toString()),
285              k);
286      BooleanQuery.Builder builder = new BooleanQuery.Builder();
287      builder.add(query, BooleanClause.Occur.SHOULD);
288      builder.add(knnQuery, BooleanClause.Occur.SHOULD);
289      return builder.build();
290    }
291    return query;
292  }
293
294  private static class QueryFieldTermExtractor extends QueryVisitor {
295    private final String field;
296    private final List<String> terms = new ArrayList<>();
297
298    QueryFieldTermExtractor(String field) {
299      this.field = field;
300    }
301
302    @Override
303    public boolean acceptField(String field) {
304      return field.equals(this.field);
305    }
306
307    @Override
308    public void consumeTerms(Query query, Term... terms) {
309      for (Term term : terms) {
310        this.terms.add(term.text());
311      }
312    }
313
314    @Override
315    public QueryVisitor getSubVisitor(BooleanClause.Occur occur, Query parent) {
316      if (occur == BooleanClause.Occur.MUST_NOT) {
317        return QueryVisitor.EMPTY_VISITOR;
318      }
319      return this;
320    }
321  }
322}