Source code

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.StoredFields;
035import org.apache.lucene.index.Term;
036import org.apache.lucene.queryparser.classic.QueryParser;
037import org.apache.lucene.search.BooleanClause;
038import org.apache.lucene.search.BooleanQuery;
039import org.apache.lucene.search.IndexSearcher;
040import org.apache.lucene.search.KnnFloatVectorQuery;
041import org.apache.lucene.search.Query;
042import org.apache.lucene.search.QueryVisitor;
043import org.apache.lucene.search.ScoreDoc;
044import org.apache.lucene.search.TopDocs;
045import org.apache.lucene.store.FSDirectory;
046import org.apache.lucene.util.IOUtils;
047
048/** Simple command-line based search demo. */
049public class SearchFiles {
050
051  private SearchFiles() {}
052
053  /** Simple command-line based search demo. */
054  public static void main(String[] args) throws Exception {
055    String usage =
056        "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.";
057    if (args.length > 0 && ("-h".equals(args[0]) || "-help".equals(args[0]))) {
058      System.out.println(usage);
059      System.exit(0);
060    }
061
062    String index = "index";
063    String field = "contents";
064    String queries = null;
065    int repeat = 0;
066    boolean raw = false;
067    int knnVectors = 0;
068    String queryString = null;
069    int hitsPerPage = 10;
070
071    for (int i = 0; i < args.length; i++) {
072      switch (args[i]) {
073        case "-index":
074          index = args[++i];
075          break;
076        case "-field":
077          field = args[++i];
078          break;
079        case "-queries":
080          queries = args[++i];
081          break;
082        case "-query":
083          queryString = args[++i];
084          break;
085        case "-repeat":
086          repeat = Integer.parseInt(args[++i]);
087          break;
088        case "-raw":
089          raw = true;
090          break;
091        case "-paging":
092          hitsPerPage = Integer.parseInt(args[++i]);
093          if (hitsPerPage <= 0) {
094            System.err.println("There must be at least 1 hit per page.");
095            System.exit(1);
096          }
097          break;
098        case "-knn_vector":
099          knnVectors = Integer.parseInt(args[++i]);
100          break;
101        default:
102          System.err.println("Unknown argument: " + args[i]);
103          System.exit(1);
104      }
105    }
106
107    DirectoryReader reader = DirectoryReader.open(FSDirectory.open(Paths.get(index)));
108    IndexSearcher searcher = new IndexSearcher(reader);
109    Analyzer analyzer = new StandardAnalyzer();
110    KnnVectorDict vectorDict = null;
111    if (knnVectors > 0) {
112      vectorDict = new KnnVectorDict(reader.directory(), IndexFiles.KNN_DICT);
113    }
114    BufferedReader in;
115    if (queries != null) {
116      in = Files.newBufferedReader(Paths.get(queries), StandardCharsets.UTF_8);
117    } else {
118      in = new BufferedReader(new InputStreamReader(System.in, StandardCharsets.UTF_8));
119    }
120    QueryParser parser = new QueryParser(field, analyzer);
121    while (true) {
122      if (queries == null && queryString == null) { // prompt the user
123        System.out.println("Enter query: ");
124      }
125
126      String line = queryString != null ? queryString : in.readLine();
127
128      if (line == null || line.length() == -1) {
129        break;
130      }
131
132      line = line.trim();
133      if (line.length() == 0) {
134        break;
135      }
136
137      Query query = parser.parse(line);
138      if (knnVectors > 0) {
139        query = addSemanticQuery(query, vectorDict, knnVectors);
140      }
141      System.out.println("Searching for: " + query.toString(field));
142
143      if (repeat > 0) { // repeat & time as benchmark
144        Date start = new Date();
145        for (int i = 0; i < repeat; i++) {
146          searcher.search(query, 100);
147        }
148        Date end = new Date();
149        System.out.println("Time: " + (end.getTime() - start.getTime()) + "ms");
150      }
151
152      doPagingSearch(in, searcher, query, hitsPerPage, raw, queries == null && queryString == null);
153
154      if (queryString != null) {
155        break;
156      }
157    }
158    IOUtils.close(vectorDict, reader);
159  }
160
161  /**
162   * This demonstrates a typical paging search scenario, where the search engine presents pages of
163   * size n to the user. The user can then go to the next page if interested in the next hits.
164   *
165   * <p>When the query is executed for the first time, then only enough results are collected to
166   * fill 5 result pages. If the user wants to page beyond this limit, then the query is executed
167   * another time and all hits are collected.
168   */
169  public static void doPagingSearch(
170      BufferedReader in,
171      IndexSearcher searcher,
172      Query query,
173      int hitsPerPage,
174      boolean raw,
175      boolean interactive)
176      throws IOException {
177
178    // Collect enough docs to show 5 pages
179    TopDocs results = searcher.search(query, 5 * hitsPerPage);
180    ScoreDoc[] hits = results.scoreDocs;
181
182    int numTotalHits = Math.toIntExact(results.totalHits.value);
183    System.out.println(numTotalHits + " total matching documents");
184
185    int start = 0;
186    int end = Math.min(numTotalHits, hitsPerPage);
187
188    while (true) {
189      if (end > hits.length) {
190        System.out.println(
191            "Only results 1 - "
192                + hits.length
193                + " of "
194                + numTotalHits
195                + " total matching documents collected.");
196        System.out.println("Collect more (y/n) ?");
197        String line = in.readLine();
198        if (line == null || line.length() == 0 || line.charAt(0) == 'n') {
199          break;
200        }
201
202        hits = searcher.search(query, numTotalHits).scoreDocs;
203      }
204
205      end = Math.min(hits.length, start + hitsPerPage);
206
207      StoredFields storedFields = searcher.storedFields();
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 = storedFields.document(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 == null || 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      KnnFloatVectorQuery knnQuery =
282          new KnnFloatVectorQuery(
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}