org.apache.lucene.search.TermInSetQuery Java Examples

/**
 * Returns a {@link DocumentQueryBuilder} for this presearcher
 */
protected DocumentQueryBuilder getQueryBuilder() {
  return new DocumentQueryBuilder() {

    Map<String, List<BytesRef>> terms = new HashMap<>();

    @Override
    public void addTerm(String field, BytesRef term) {
      List<BytesRef> t = terms.computeIfAbsent(field, f -> new ArrayList<>());
      t.add(term);
    }

    @Override
    public Query build() {
      BooleanQuery.Builder builder = new BooleanQuery.Builder();
      for (Map.Entry<String, List<BytesRef>> entry : terms.entrySet()) {
        builder.add(new TermInSetQuery(entry.getKey(), entry.getValue()), BooleanClause.Occur.SHOULD);
      }
      return builder.build();
    }
  };
}
 

@Override
public Query build() {
  Map<String, BytesRef[]> collectedTerms = new HashMap<>();
  for (Map.Entry<String, BytesRefHash> entry : terms.entrySet()) {
    collectedTerms.put(entry.getKey(), convertHash(entry.getValue()));
  }
  BooleanQuery.Builder parent = new BooleanQuery.Builder();
  for (int i = 0; i < passes; i++) {
    BooleanQuery.Builder child = new BooleanQuery.Builder();
    for (String field : terms.keySet()) {
      child.add(new TermInSetQuery(field(field, i), collectedTerms.get(field)), BooleanClause.Occur.SHOULD);
    }
    parent.add(child.build(), BooleanClause.Occur.MUST);
  }
  return parent.build();
}
 

@Test
public void testAnyEqArrayLiteral() throws Exception {
    Query query = convert("d = any([-1.5, 0.0, 1.5])");
    assertThat(query, instanceOf(PointInSetQuery.class));

    query = convert("_id in ('test','test2')");
    assertThat(query, instanceOf(TermInSetQuery.class));

    query = convert("_id in (1, 2)");
    assertThat(query, instanceOf(TermInSetQuery.class));

    query = convert("_id = any (['test','test2'])");
    assertThat(query, instanceOf(TermInSetQuery.class));

    query = convert("_id = any ([1, 2])");
    assertThat(query, instanceOf(TermInSetQuery.class));
}
 

private Query getGroupQuery(String fname,
                         FieldType ft,
                         int size,
                         LongHashSet groupSet) {

  BytesRef[] bytesRefs = new BytesRef[size];
  int index = -1;
  BytesRefBuilder term = new BytesRefBuilder();
  Iterator<LongCursor> it = groupSet.iterator();

  while (it.hasNext()) {
    LongCursor cursor = it.next();
    String stringVal = numericToString(ft, cursor.value);
    ft.readableToIndexed(stringVal, term);
    bytesRefs[++index] = term.toBytesRef();
  }

  return new TermInSetQuery(fname, bytesRefs);
}
 

@Override
public Query termsQuery(List<?> values, QueryShardContext context) {
    failIfNotIndexed();
    BytesRef[] bytesRefs = new BytesRef[values.size()];
    for (int i = 0; i < bytesRefs.length; i++) {
        bytesRefs[i] = indexedValueForSearch(values.get(i));
    }
    return new TermInSetQuery(name(), bytesRefs);
}
 

@Test
public void testIdQuery() throws Exception {
    Query query = convert("_id = 'i1'");
    assertThat(query, instanceOf(TermInSetQuery.class));

    query = convert("_id = 1");
    assertThat(query, instanceOf(TermInSetQuery.class));
}
 

@Override
public Query termsQuery(List<?> values, QueryShardContext context) {
    failIfNotIndexed();
    BytesRef[] bytesRefs = new BytesRef[values.size()];
    for (int i = 0; i < bytesRefs.length; i++) {
        bytesRefs[i] = indexedValueForSearch(values.get(i));
    }
    return new TermInSetQuery(name(), bytesRefs);
}
 

@Override
public Query termsQuery(List<?> values, QueryShardContext context) {
    failIfNotIndexed();
    BytesRef[] bytesRefs = new BytesRef[values.size()];
    for (int i = 0; i < bytesRefs.length; i++) {
        Object idObject = values.get(i);
        if (idObject instanceof BytesRef) {
            idObject = ((BytesRef) idObject).utf8ToString();
        }
        bytesRefs[i] = Uid.encodeId(idObject.toString());
    }
    return new TermInSetQuery(name(), bytesRefs);
}
 

@Override
public Query termsQuery(List<?> values, QueryShardContext context) {
    failIfNotIndexed();
    BytesRef[] bytesRefs = new BytesRef[values.size()];
    for (int i = 0; i < bytesRefs.length; i++) {
        bytesRefs[i] = indexedValueForSearch(values.get(i));
    }
    return new TermInSetQuery(name(), bytesRefs);
}
 

private Query getGroupQuery(String fname,
                            int size,
                            IntObjectHashMap<BytesRef> ordBytes) {
  BytesRef[] bytesRefs = new BytesRef[size];
  int index = -1;
  Iterator<IntObjectCursor<BytesRef>>it = ordBytes.iterator();
  while (it.hasNext()) {
    IntObjectCursor<BytesRef> cursor = it.next();
    bytesRefs[++index] = cursor.value;
  }
  return new TermInSetQuery(fname, bytesRefs);
}
 

@Override
public Query getResultQuery(SchemaField matchField, boolean useAutomaton) {
  if (collectorTerms == null || collectorTerms.size() == 0) {
    // return null if there are no terms (edges) to traverse.
    return null;
  } else {
    // Create a query
    Query q = null;

    // TODO: see if we should dynamically select this based on the frontier size.
    if (useAutomaton) {
      // build an automaton based query for the frontier.
      Automaton autn = buildAutomaton(collectorTerms);
      AutomatonQuery autnQuery = new AutomatonQuery(new Term(matchField.getName()), autn);
      q = autnQuery;
    } else {
      List<BytesRef> termList = new ArrayList<>(collectorTerms.size());
      for (int i = 0; i < collectorTerms.size(); i++) {
        BytesRef ref = new BytesRef();
        collectorTerms.get(i, ref);
        termList.add(ref);
      }
      q = (matchField.hasDocValues() && !matchField.indexed())
              ? new DocValuesTermsQuery(matchField.getName(), termList)
              : new TermInSetQuery(matchField.getName(), termList);
    }

    return q;
  }
}
 

@Override
@SuppressWarnings("unchecked")
Query makeQuery(String fname, Iterator<BytesRef> it) {
  return new TermInSetQuery(fname, IteratorUtils.toList(it));
}
 

@Override
public Query makeQuery(SpatialArgs args) {
  final SpatialOperation op = args.getOperation();
  if (op != SpatialOperation.Intersects)
    throw new UnsupportedSpatialOperation(op);

  Shape shape = args.getShape();
  int detailLevel = grid.getLevelForDistance(args.resolveDistErr(ctx, distErrPct));

  //--get a List of BytesRef for each term we want (no parents, no leaf bytes))
  final int GUESS_NUM_TERMS;
  if (shape instanceof Point)
    GUESS_NUM_TERMS = detailLevel;//perfect guess
  else
    GUESS_NUM_TERMS = 4096;//should this be a method on SpatialPrefixTree?

  BytesRefBuilder masterBytes = new BytesRefBuilder();//shared byte array for all terms
  List<BytesRef> terms = new ArrayList<>(GUESS_NUM_TERMS);

  CellIterator cells = grid.getTreeCellIterator(shape, detailLevel);
  while (cells.hasNext()) {
    Cell cell = cells.next();
    if (!cell.isLeaf())
      continue;
    BytesRef term = cell.getTokenBytesNoLeaf(null);//null because we want a new BytesRef
    //We copy out the bytes because it may be re-used across the iteration. This also gives us the opportunity
    // to use one contiguous block of memory for the bytes of all terms we need.
    masterBytes.grow(masterBytes.length() + term.length);
    masterBytes.append(term);
    term.bytes = null;//don't need; will reset later
    term.offset = masterBytes.length() - term.length;
    terms.add(term);
  }
  //doing this now because if we did earlier, it's possible the bytes needed to grow()
  for (BytesRef byteRef : terms) {
    byteRef.bytes = masterBytes.bytes();
  }
  //unfortunately TermsQuery will needlessly sort & dedupe
  //TODO an automatonQuery might be faster?
  return new TermInSetQuery(getFieldName(), terms);
}
 

@Test
public void testWhereStringRefInSetLiteralIsConvertedToTermsQuery() throws Exception {
    Query query = convert("name in ('foo', 'bar')");
    assertThat(query, instanceOf(TermInSetQuery.class));
}
 

public void testQueryBuilder() throws IOException {

    Presearcher presearcher = createPresearcher();

    IndexWriterConfig iwc = new IndexWriterConfig(new KeywordAnalyzer());
    Directory dir = new ByteBuffersDirectory();
    IndexWriter writer = new IndexWriter(dir, iwc);
    MonitorConfiguration config = new MonitorConfiguration(){
      @Override
      public IndexWriter buildIndexWriter() {
        return writer;
      }
    };

    try (Monitor monitor = new Monitor(ANALYZER, presearcher, config)) {

      monitor.register(new MonitorQuery("1", parse("f:test")));

      try (IndexReader reader = DirectoryReader.open(writer, false, false)) {

        MemoryIndex mindex = new MemoryIndex();
        mindex.addField("f", "this is a test document", WHITESPACE);
        mindex.addField("g", "#######", ANALYZER); // analyzes away to empty field
        LeafReader docsReader = (LeafReader) mindex.createSearcher().getIndexReader();

        QueryIndex.QueryTermFilter termFilter = new QueryIndex.QueryTermFilter(reader);

        BooleanQuery q = (BooleanQuery) presearcher.buildQuery(docsReader, termFilter);
        BooleanQuery expected = new BooleanQuery.Builder()
            .add(should(new BooleanQuery.Builder()
                .add(should(new TermInSetQuery("f", new BytesRef("test")))).build()))
            .add(should(new TermQuery(new Term("__anytokenfield", "__ANYTOKEN__"))))
            .build();

        assertEquals(expected, q);

      }

    }

  }
 

public void testQueryBuilder() throws IOException {

    IndexWriterConfig iwc = new IndexWriterConfig(new KeywordAnalyzer());
    Presearcher presearcher = createPresearcher();

    Directory dir = new ByteBuffersDirectory();
    IndexWriter writer = new IndexWriter(dir, iwc);
    MonitorConfiguration config = new MonitorConfiguration(){
      @Override
      public IndexWriter buildIndexWriter() {
        return writer;
      }
    };
    try (Monitor monitor = new Monitor(ANALYZER, presearcher, config)) {

      monitor.register(new MonitorQuery("1", parse("f:test")));

      try (IndexReader reader = DirectoryReader.open(writer, false, false)) {

        MemoryIndex mindex = new MemoryIndex();
        mindex.addField("f", "this is a test document", WHITESPACE);
        LeafReader docsReader = (LeafReader) mindex.createSearcher().getIndexReader();

        QueryIndex.QueryTermFilter termFilter = new QueryIndex.QueryTermFilter(reader);

        BooleanQuery q = (BooleanQuery) presearcher.buildQuery(docsReader, termFilter);
        BooleanQuery expected = new BooleanQuery.Builder()
            .add(should(new BooleanQuery.Builder()
                .add(must(new BooleanQuery.Builder().add(should(new TermInSetQuery("f_0", new BytesRef("test")))).build()))
                .add(must(new BooleanQuery.Builder().add(should(new TermInSetQuery("f_1", new BytesRef("test")))).build()))
                .add(must(new BooleanQuery.Builder().add(should(new TermInSetQuery("f_2", new BytesRef("test")))).build()))
                .add(must(new BooleanQuery.Builder().add(should(new TermInSetQuery("f_3", new BytesRef("test")))).build()))
                .build()))
            .add(should(new TermQuery(new Term("__anytokenfield", "__ANYTOKEN__"))))
            .build();

        assertEquals(expected, q);
      }

    }

  }
 

@Test
public void testWhereInIsOptimized() throws Exception {
    Query query = convert("name in ('foo', 'bar')");
    assertThat(query, instanceOf(TermInSetQuery.class));
    assertThat(query.toString(), is("name:(bar foo)"));
}