Java 类org.apache.lucene.search.TermInSetQuery 实例源码

@Override
protected Query doToQuery(QueryShardContext context) throws IOException {
    if (termsLookup != null) {
        throw new UnsupportedOperationException("query must be rewritten first");
    }
    if (values == null || values.isEmpty()) {
        return Queries.newMatchNoDocsQuery("No terms supplied for \"" + getName() + "\" query.");
    }
    MappedFieldType fieldType = context.fieldMapper(fieldName);

    if (fieldType != null) {
        return fieldType.termsQuery(values, context);
    } else {
        BytesRef[] filterValues = new BytesRef[values.size()];
        for (int i = 0; i < filterValues.length; i++) {
            filterValues[i] = BytesRefs.toBytesRef(values.get(i));
        }
        return new TermInSetQuery(fieldName, filterValues);
    }
}

Query createCandidateQuery(IndexReader indexReader) throws IOException {
    List<BytesRef> extractedTerms = new ArrayList<>();
    LeafReader reader = indexReader.leaves().get(0).reader();
    Fields fields = reader.fields();
    for (String field : fields) {
        Terms terms = fields.terms(field);
        if (terms == null) {
            continue;
        }

        BytesRef fieldBr = new BytesRef(field);
        TermsEnum tenum = terms.iterator();
        for (BytesRef term = tenum.next(); term != null; term = tenum.next()) {
            BytesRefBuilder builder = new BytesRefBuilder();
            builder.append(fieldBr);
            builder.append(FIELD_VALUE_SEPARATOR);
            builder.append(term);
            extractedTerms.add(builder.toBytesRef());
        }
    }
    Query extractionSuccess = new TermInSetQuery(queryTermsField.name(), extractedTerms);
    // include extractionResultField:failed, because docs with this term have no extractedTermsField
    // and otherwise we would fail to return these docs. Docs that failed query term extraction
    // always need to be verified by MemoryIndex:
    Query extractionFailure = new TermQuery(new Term(extractionResultField.name(), EXTRACTION_FAILED));

    return new BooleanQuery.Builder()
            .add(extractionSuccess, Occur.SHOULD)
            .add(extractionFailure, Occur.SHOULD)
            .build();
}

public void testExtractQueryMetadata_termsQuery() {
    TermInSetQuery termsQuery = new TermInSetQuery("_field", new BytesRef("_term1"), new BytesRef("_term2"));
    Result result = analyze(termsQuery);
    assertThat(result.verified, is(true));
    List<Term> terms = new ArrayList<>(result.terms);
    Collections.sort(terms);
    assertThat(terms.size(), equalTo(2));
    assertThat(terms.get(0).field(), equalTo("_field"));
    assertThat(terms.get(0).text(), equalTo("_term1"));
    assertThat(terms.get(1).field(), equalTo("_field"));
    assertThat(terms.get(1).text(), equalTo("_term2"));
}

public void testCreateCandidateQuery() throws Exception {
    addQueryMapping();

    MemoryIndex memoryIndex = new MemoryIndex(false);
    memoryIndex.addField("field1", "the quick brown fox jumps over the lazy dog", new WhitespaceAnalyzer());
    memoryIndex.addField("field2", "some more text", new WhitespaceAnalyzer());
    memoryIndex.addField("_field3", "unhide me", new WhitespaceAnalyzer());
    memoryIndex.addField("field4", "123", new WhitespaceAnalyzer());
    memoryIndex.addField(new LongPoint("number_field", 10L), new WhitespaceAnalyzer());

    IndexReader indexReader = memoryIndex.createSearcher().getIndexReader();

    BooleanQuery candidateQuery = (BooleanQuery) fieldType.createCandidateQuery(indexReader);
    assertEquals(2, candidateQuery.clauses().size());
    assertEquals(Occur.SHOULD, candidateQuery.clauses().get(0).getOccur());
    TermInSetQuery termsQuery = (TermInSetQuery) candidateQuery.clauses().get(0).getQuery();

    PrefixCodedTerms terms = termsQuery.getTermData();
    assertThat(terms.size(), equalTo(14L));
    PrefixCodedTerms.TermIterator termIterator = terms.iterator();
    assertTermIterator(termIterator, "_field3\u0000me", fieldType.queryTermsField.name());
    assertTermIterator(termIterator, "_field3\u0000unhide", fieldType.queryTermsField.name());
    assertTermIterator(termIterator, "field1\u0000brown", fieldType.queryTermsField.name());
    assertTermIterator(termIterator, "field1\u0000dog", fieldType.queryTermsField.name());
    assertTermIterator(termIterator, "field1\u0000fox", fieldType.queryTermsField.name());
    assertTermIterator(termIterator, "field1\u0000jumps", fieldType.queryTermsField.name());
    assertTermIterator(termIterator, "field1\u0000lazy", fieldType.queryTermsField.name());
    assertTermIterator(termIterator, "field1\u0000over", fieldType.queryTermsField.name());
    assertTermIterator(termIterator, "field1\u0000quick", fieldType.queryTermsField.name());
    assertTermIterator(termIterator, "field1\u0000the", fieldType.queryTermsField.name());
    assertTermIterator(termIterator, "field2\u0000more", fieldType.queryTermsField.name());
    assertTermIterator(termIterator, "field2\u0000some", fieldType.queryTermsField.name());
    assertTermIterator(termIterator, "field2\u0000text", fieldType.queryTermsField.name());
    assertTermIterator(termIterator, "field4\u0000123", fieldType.queryTermsField.name());

    assertEquals(Occur.SHOULD, candidateQuery.clauses().get(1).getOccur());
    assertEquals(new TermQuery(new Term(fieldType.extractionResultField.name(), EXTRACTION_FAILED)),
            candidateQuery.clauses().get(1).getQuery());
}

private static void handleExclude(BooleanQuery.Builder boolQuery, Item[] likeItems) {
    // artificial docs get assigned a random id and should be disregarded
    List<BytesRef> uids = new ArrayList<>();
    for (Item item : likeItems) {
        if (item.doc() != null) {
            continue;
        }
        uids.add(createUidAsBytes(item.type(), item.id()));
    }
    if (!uids.isEmpty()) {
        TermInSetQuery query = new TermInSetQuery(UidFieldMapper.NAME, uids.toArray(new BytesRef[uids.size()]));
        boolQuery.add(query, BooleanClause.Occur.MUST_NOT);
    }
}

@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++) {
        bytesRefs[i] = indexedValueForSearch(values.get(i));
    }
    return new TermInSetQuery(name(), bytesRefs);
}

public void testParentChild() throws IOException {
    Directory directory = newDirectory();
    RandomIndexWriter indexWriter = new RandomIndexWriter(random(), directory);

    final Map<String, Tuple<Integer, Integer>> expectedParentChildRelations = setupIndex(indexWriter);
    indexWriter.close();

    IndexReader indexReader = ElasticsearchDirectoryReader.wrap(DirectoryReader.open(directory),
            new ShardId(new Index("foo", "_na_"), 1));
    // TODO set "maybeWrap" to true for IndexSearcher once #23338 is resolved
    IndexSearcher indexSearcher = newSearcher(indexReader, false, true);

    testCase(new MatchAllDocsQuery(), indexSearcher, child -> {
        int expectedTotalChildren = 0;
        int expectedMinValue = Integer.MAX_VALUE;
        for (Tuple<Integer, Integer> expectedValues : expectedParentChildRelations.values()) {
            expectedTotalChildren += expectedValues.v1();
            expectedMinValue = Math.min(expectedMinValue, expectedValues.v2());
        }
        assertEquals(expectedTotalChildren, child.getDocCount());
        assertEquals(expectedMinValue, ((InternalMin) child.getAggregations().get("in_child")).getValue(), Double.MIN_VALUE);
    });

    for (String parent : expectedParentChildRelations.keySet()) {
        testCase(new TermInSetQuery(UidFieldMapper.NAME, new BytesRef(Uid.createUid(PARENT_TYPE, parent))), indexSearcher, child -> {
            assertEquals((long) expectedParentChildRelations.get(parent).v1(), child.getDocCount());
            assertEquals(expectedParentChildRelations.get(parent).v2(),
                    ((InternalMin) child.getAggregations().get("in_child")).getValue(), Double.MIN_VALUE);
        });
    }
    indexReader.close();
    directory.close();
}

static void assertLateParsingQuery(Query query, String type, String id) throws IOException {
    assertThat(query, instanceOf(HasChildQueryBuilder.LateParsingQuery.class));
    HasChildQueryBuilder.LateParsingQuery lateParsingQuery = (HasChildQueryBuilder.LateParsingQuery) query;
    assertThat(lateParsingQuery.getInnerQuery(), instanceOf(BooleanQuery.class));
    BooleanQuery booleanQuery = (BooleanQuery) lateParsingQuery.getInnerQuery();
    assertThat(booleanQuery.clauses().size(), equalTo(2));
    //check the inner ids query, we have to call rewrite to get to check the type it's executed against
    assertThat(booleanQuery.clauses().get(0).getOccur(), equalTo(BooleanClause.Occur.MUST));
    assertThat(booleanQuery.clauses().get(0).getQuery(), instanceOf(TermInSetQuery.class));
    TermInSetQuery termsQuery = (TermInSetQuery) booleanQuery.clauses().get(0).getQuery();
    Query rewrittenTermsQuery = termsQuery.rewrite(null);
    assertThat(rewrittenTermsQuery, instanceOf(ConstantScoreQuery.class));
    ConstantScoreQuery constantScoreQuery = (ConstantScoreQuery) rewrittenTermsQuery;
    assertThat(constantScoreQuery.getQuery(), instanceOf(BooleanQuery.class));
    BooleanQuery booleanTermsQuery = (BooleanQuery) constantScoreQuery.getQuery();
    assertThat(booleanTermsQuery.clauses().toString(), booleanTermsQuery.clauses().size(), equalTo(1));
    assertThat(booleanTermsQuery.clauses().get(0).getOccur(), equalTo(BooleanClause.Occur.SHOULD));
    assertThat(booleanTermsQuery.clauses().get(0).getQuery(), instanceOf(TermQuery.class));
    TermQuery termQuery = (TermQuery) booleanTermsQuery.clauses().get(0).getQuery();
    assertThat(termQuery.getTerm().field(), equalTo(UidFieldMapper.NAME));
    //we want to make sure that the inner ids query gets executed against the child type rather than the main type we initially set to the context
    BytesRef[] ids = Uid.createUidsForTypesAndIds(Collections.singletonList(type), Collections.singletonList(id));
    assertThat(termQuery.getTerm().bytes(), equalTo(ids[0]));
    //check the type filter
    assertThat(booleanQuery.clauses().get(1).getOccur(), equalTo(BooleanClause.Occur.FILTER));
    assertEquals(new TypeFieldMapper.TypesQuery(new BytesRef(type)), booleanQuery.clauses().get(1).getQuery());
}

@Override
protected void doAssertLuceneQuery(IdsQueryBuilder queryBuilder, Query query, SearchContext context) throws IOException {
    if (queryBuilder.ids().size() == 0) {
        assertThat(query, instanceOf(MatchNoDocsQuery.class));
    } else {
        assertThat(query, instanceOf(TermInSetQuery.class));
    }
}

public void testTermsQuery() {
    MappedFieldType ft = createDefaultFieldType();
    ft.setName("field");
    ft.setIndexOptions(IndexOptions.DOCS);
    List<BytesRef> terms = new ArrayList<>();
    terms.add(new BytesRef("foo"));
    terms.add(new BytesRef("bar"));
    assertEquals(new TermInSetQuery("field", terms),
            ft.termsQuery(Arrays.asList("foo", "bar"), null));

    ft.setIndexOptions(IndexOptions.NONE);
    IllegalArgumentException e = expectThrows(IllegalArgumentException.class,
            () -> ft.termsQuery(Arrays.asList("foo", "bar"), null));
    assertEquals("Cannot search on field [field] since it is not indexed.", e.getMessage());
}

public void testTermsQuery() {
    MappedFieldType ft = createDefaultFieldType();
    ft.setName("field");
    ft.setIndexOptions(IndexOptions.DOCS);
    List<BytesRef> terms = new ArrayList<>();
    terms.add(new BytesRef("foo"));
    terms.add(new BytesRef("bar"));
    assertEquals(new TermInSetQuery("field", terms),
            ft.termsQuery(Arrays.asList("foo", "bar"), null));

    ft.setIndexOptions(IndexOptions.NONE);
    IllegalArgumentException e = expectThrows(IllegalArgumentException.class,
            () -> ft.termsQuery(Arrays.asList("foo", "bar"), null));
    assertEquals("Cannot search on field [field] since it is not indexed.", e.getMessage());
}

@Override
public <P extends ParaObject> List<P> findTermInList(String appid, String type, String field,
        List<?> terms, Pager... pager) {
    if (StringUtils.isBlank(field) || terms == null) {
        return Collections.emptyList();
    }
    ArrayList<BytesRef> termsList = new ArrayList<>();
    for (Object term : terms) {
        termsList.add(new BytesRef(term.toString()));
    }
    Query query = new TermInSetQuery(field, termsList);
    return searchQuery(dao, appid, type, query, pager);
}

@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 termQuery(Object value, @Nullable QueryShardContext context) {
    final BytesRef[] uids = Uid.createUidsForTypesAndId(context.queryTypes(), value);
    return new TermInSetQuery(UidFieldMapper.NAME, uids);
}

@Override
public Query termsQuery(List values, @Nullable QueryShardContext context) {
    return new TermInSetQuery(UidFieldMapper.NAME, Uid.createUidsForTypesAndIds(context.queryTypes(), values));
}

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