Example usage for org.apache.lucene.index IndexReader leaves

List of usage examples for org.apache.lucene.index IndexReader leaves

Introduction

In this page you can find the example usage for org.apache.lucene.index IndexReader leaves.

Prototype

public final List<LeafReaderContext> leaves() 

Source Link

Document

Returns the reader's leaves, or itself if this reader is atomic.

Usage

From source file:org.elasticsearch.search.aggregations.bucket.children.ParentToChildrenAggregator.java

License:Apache License

@Override
protected void doPostCollection() throws IOException {
    IndexReader indexReader = context().searchContext().searcher().getIndexReader();
    for (LeafReaderContext ctx : indexReader.leaves()) {
        Scorer childDocsScorer = childFilter.scorer(ctx);
        if (childDocsScorer == null) {
            continue;
        }//from ww w.j ava2  s. c o  m
        DocIdSetIterator childDocsIter = childDocsScorer.iterator();

        final LeafBucketCollector sub = collectableSubAggregators.getLeafCollector(ctx);
        final SortedDocValues globalOrdinals = valuesSource.globalOrdinalsValues(parentType, ctx);

        // Set the scorer, since we now replay only the child docIds
        sub.setScorer(ConstantScorer.create(childDocsIter, null, 1f));

        final Bits liveDocs = ctx.reader().getLiveDocs();
        for (int docId = childDocsIter.nextDoc(); docId != DocIdSetIterator.NO_MORE_DOCS; docId = childDocsIter
                .nextDoc()) {
            if (liveDocs != null && liveDocs.get(docId) == false) {
                continue;
            }
            long globalOrdinal = globalOrdinals.getOrd(docId);
            if (globalOrdinal != -1) {
                long bucketOrd = parentOrdToBuckets.get(globalOrdinal);
                if (bucketOrd != -1) {
                    collectBucket(sub, docId, bucketOrd);
                    if (multipleBucketsPerParentOrd) {
                        long[] otherBucketOrds = parentOrdToOtherBuckets.get(globalOrdinal);
                        if (otherBucketOrds != null) {
                            for (long otherBucketOrd : otherBucketOrds) {
                                collectBucket(sub, docId, otherBucketOrd);
                            }
                        }
                    }
                }
            }
        }
    }
}

From source file:org.elasticsearch.search.aggregations.bucket.composite.CompositeValuesCollectorQueueTests.java

License:Apache License

private void testRandomCase(boolean forceMerge, ClassAndName... types) throws IOException {
    final BigArrays bigArrays = BigArrays.NON_RECYCLING_INSTANCE;
    int numDocs = randomIntBetween(50, 100);
    List<Comparable<?>[]> possibleValues = new ArrayList<>();
    for (ClassAndName type : types) {
        int numValues = randomIntBetween(1, numDocs * 2);
        Comparable<?>[] values = new Comparable[numValues];
        if (type.clazz == Long.class) {
            for (int i = 0; i < numValues; i++) {
                values[i] = randomLong();
            }/* w ww . ja v  a  2  s.  c  o  m*/
        } else if (type.clazz == Double.class) {
            for (int i = 0; i < numValues; i++) {
                values[i] = randomDouble();
            }
        } else if (type.clazz == BytesRef.class) {
            for (int i = 0; i < numValues; i++) {
                values[i] = new BytesRef(randomAlphaOfLengthBetween(5, 50));
            }
        } else {
            assert (false);
        }
        possibleValues.add(values);
    }

    Set<CompositeKey> keys = new HashSet<>();
    try (Directory directory = newDirectory()) {
        try (RandomIndexWriter indexWriter = new RandomIndexWriter(random(), directory,
                new KeywordAnalyzer())) {
            for (int i = 0; i < numDocs; i++) {
                Document document = new Document();
                List<List<Comparable<?>>> docValues = new ArrayList<>();
                boolean hasAllField = true;
                for (int j = 0; j < types.length; j++) {
                    int numValues = randomIntBetween(0, 5);
                    if (numValues == 0) {
                        hasAllField = false;
                    }
                    List<Comparable<?>> values = new ArrayList<>();
                    for (int k = 0; k < numValues; k++) {
                        values.add(
                                possibleValues.get(j)[randomIntBetween(0, possibleValues.get(j).length - 1)]);
                        if (types[j].clazz == Long.class) {
                            long value = (Long) values.get(k);
                            document.add(new SortedNumericDocValuesField(types[j].fieldType.name(), value));
                            document.add(new LongPoint(types[j].fieldType.name(), value));
                        } else if (types[j].clazz == Double.class) {
                            document.add(new SortedNumericDocValuesField(types[j].fieldType.name(),
                                    NumericUtils.doubleToSortableLong((Double) values.get(k))));
                        } else if (types[j].clazz == BytesRef.class) {
                            BytesRef value = (BytesRef) values.get(k);
                            document.add(new SortedSetDocValuesField(types[j].fieldType.name(),
                                    (BytesRef) values.get(k)));
                            document.add(new TextField(types[j].fieldType.name(), value.utf8ToString(),
                                    Field.Store.NO));
                        } else {
                            assert (false);
                        }
                    }
                    docValues.add(values);
                }
                if (hasAllField) {
                    List<CompositeKey> comb = createListCombinations(docValues);
                    keys.addAll(comb);
                }
                indexWriter.addDocument(document);
            }
            if (forceMerge) {
                indexWriter.forceMerge(1);
            }
        }
        IndexReader reader = DirectoryReader.open(directory);
        int size = randomIntBetween(1, keys.size());
        SingleDimensionValuesSource<?>[] sources = new SingleDimensionValuesSource[types.length];
        for (int i = 0; i < types.length; i++) {
            final MappedFieldType fieldType = types[i].fieldType;
            if (types[i].clazz == Long.class) {
                sources[i] = new LongValuesSource(bigArrays, fieldType,
                        context -> context.reader().getSortedNumericDocValues(fieldType.name()), value -> value,
                        DocValueFormat.RAW, size, 1);
            } else if (types[i].clazz == Double.class) {
                sources[i] = new DoubleValuesSource(bigArrays, fieldType,
                        context -> FieldData.sortableLongBitsToDoubles(
                                context.reader().getSortedNumericDocValues(fieldType.name())),
                        size, 1);
            } else if (types[i].clazz == BytesRef.class) {
                if (forceMerge) {
                    // we don't create global ordinals but we test this mode when the reader has a single segment
                    // since ordinals are global in this case.
                    sources[i] = new GlobalOrdinalValuesSource(bigArrays, fieldType,
                            context -> context.reader().getSortedSetDocValues(fieldType.name()), size, 1);
                } else {
                    sources[i] = new BinaryValuesSource(fieldType,
                            context -> FieldData
                                    .toString(context.reader().getSortedSetDocValues(fieldType.name())),
                            size, 1);
                }
            } else {
                assert (false);
            }
        }
        CompositeKey[] expected = keys.toArray(new CompositeKey[0]);
        Arrays.sort(expected, (a, b) -> compareKey(a, b));
        CompositeValuesCollectorQueue queue = new CompositeValuesCollectorQueue(sources, size);
        final SortedDocsProducer docsProducer = sources[0].createSortedDocsProducerOrNull(reader,
                new MatchAllDocsQuery());
        for (boolean withProducer : new boolean[] { true, false }) {
            if (withProducer && docsProducer == null) {
                continue;
            }
            int pos = 0;
            CompositeKey last = null;
            while (pos < size) {
                queue.clear();
                if (last != null) {
                    queue.setAfter(last.values());
                }

                for (LeafReaderContext leafReaderContext : reader.leaves()) {
                    final LeafBucketCollector leafCollector = new LeafBucketCollector() {
                        @Override
                        public void collect(int doc, long bucket) throws IOException {
                            queue.addIfCompetitive();
                        }
                    };
                    if (withProducer) {
                        assertEquals(DocIdSet.EMPTY, docsProducer.processLeaf(new MatchAllDocsQuery(), queue,
                                leafReaderContext, false));
                    } else {
                        final LeafBucketCollector queueCollector = queue.getLeafCollector(leafReaderContext,
                                leafCollector);
                        final Bits liveDocs = leafReaderContext.reader().getLiveDocs();
                        for (int i = 0; i < leafReaderContext.reader().maxDoc(); i++) {
                            if (liveDocs == null || liveDocs.get(i)) {
                                queueCollector.collect(i);
                            }
                        }
                    }
                }
                assertEquals(size, Math.min(queue.size(), expected.length - pos));
                int ptr = 0;
                for (int slot : queue.getSortedSlot()) {
                    CompositeKey key = queue.toCompositeKey(slot);
                    assertThat(key, equalTo(expected[ptr++]));
                    last = key;
                }
                pos += queue.size();
            }
        }
        reader.close();
    }
}

From source file:org.elasticsearch.search.aggregations.bucket.composite.CompositeValuesSourceBuilder.java

License:Apache License

private static boolean isSingleValued(IndexReader reader, SortField field) throws IOException {
    SortField.Type type = IndexSortConfig.getSortFieldType(field);
    for (LeafReaderContext context : reader.leaves()) {
        if (type == SortField.Type.STRING) {
            final SortedSetDocValues values = DocValues.getSortedSet(context.reader(), field.getField());
            if (values.cost() > 0 && DocValues.unwrapSingleton(values) == null) {
                return false;
            }//from   w  w  w. j av  a2s . co  m
        } else {
            final SortedNumericDocValues values = DocValues.getSortedNumeric(context.reader(),
                    field.getField());
            if (values.cost() > 0 && DocValues.unwrapSingleton(values) == null) {
                return false;
            }
        }
    }
    return true;
}

From source file:org.elasticsearch.search.aggregations.bucket.histogram.DateHistogramAggregationBuilder.java

License:Apache License

DateTimeZone rewriteTimeZone(QueryShardContext context) throws IOException {
    final DateTimeZone tz = timeZone();
    if (field() != null && tz != null && tz.isFixed() == false && field() != null && script() == null) {
        final MappedFieldType ft = context.fieldMapper(field());
        final IndexReader reader = context.getIndexReader();
        if (ft != null && reader != null) {
            Long anyInstant = null;
            final IndexNumericFieldData fieldData = context.getForField(ft);
            for (LeafReaderContext ctx : reader.leaves()) {
                AtomicNumericFieldData leafFD = fieldData.load(ctx);
                SortedNumericDocValues values = leafFD.getLongValues();
                if (values.nextDoc() != DocIdSetIterator.NO_MORE_DOCS) {
                    anyInstant = values.nextValue();
                    break;
                }// w w  w  .  j  a  va  2  s  . c o  m
            }

            if (anyInstant != null) {
                final long prevTransition = tz.previousTransition(anyInstant);
                final long nextTransition = tz.nextTransition(anyInstant);

                // We need all not only values but also rounded values to be within
                // [prevTransition, nextTransition].
                final long low;
                DateTimeUnit intervalAsUnit = getIntervalAsDateTimeUnit();
                if (intervalAsUnit != null) {
                    final DateTimeField dateTimeField = intervalAsUnit.field(tz);
                    low = dateTimeField.roundCeiling(prevTransition);
                } else {
                    final TimeValue intervalAsMillis = getIntervalAsTimeValue();
                    low = Math.addExact(prevTransition, intervalAsMillis.millis());
                }
                // rounding rounds down, so 'nextTransition' is a good upper bound
                final long high = nextTransition;

                if (ft.isFieldWithinQuery(reader, low, high, true, false, DateTimeZone.UTC, EPOCH_MILLIS_PARSER,
                        context) == Relation.WITHIN) {
                    // All values in this reader have the same offset despite daylight saving times.
                    // This is very common for location-based timezones such as Europe/Paris in
                    // combination with time-based indices.
                    return DateTimeZone.forOffsetMillis(tz.getOffset(anyInstant));
                }
            }
        }
    }
    return tz;
}

From source file:org.elasticsearch.search.fetch.innerhits.NestedChildrenFilterTest.java

License:Apache License

@Test
public void testNestedChildrenFilter() throws Exception {
    int numParentDocs = scaledRandomIntBetween(0, 32);
    int maxChildDocsPerParent = scaledRandomIntBetween(8, 16);

    Directory dir = newDirectory();/* ww w .  ja  v a 2  s  .  c o m*/
    RandomIndexWriter writer = new RandomIndexWriter(random(), dir);
    for (int i = 0; i < numParentDocs; i++) {
        int numChildDocs = scaledRandomIntBetween(0, maxChildDocsPerParent);
        List<Document> docs = new ArrayList<>(numChildDocs + 1);
        for (int j = 0; j < numChildDocs; j++) {
            Document childDoc = new Document();
            childDoc.add(new StringField("type", "child", Field.Store.NO));
            docs.add(childDoc);
        }

        Document parenDoc = new Document();
        parenDoc.add(new StringField("type", "parent", Field.Store.NO));
        parenDoc.add(new IntField("num_child_docs", numChildDocs, Field.Store.YES));
        docs.add(parenDoc);
        writer.addDocuments(docs);
    }

    IndexReader reader = writer.getReader();
    writer.close();

    IndexSearcher searcher = new IndexSearcher(reader);
    FetchSubPhase.HitContext hitContext = new FetchSubPhase.HitContext();
    FixedBitSetFilterCache fixedBitSetFilterCache = new FixedBitSetFilterCache(new Index("test"),
            ImmutableSettings.EMPTY);
    FixedBitSetFilter parentFilter = fixedBitSetFilterCache
            .getFixedBitSetFilter(new TermFilter(new Term("type", "parent")));
    Filter childFilter = new TermFilter(new Term("type", "child"));
    int checkedParents = 0;
    for (AtomicReaderContext leaf : reader.leaves()) {
        DocIdSetIterator parents = parentFilter.getDocIdSet(leaf, null).iterator();
        for (int parentDoc = parents.nextDoc(); parentDoc != DocIdSetIterator.NO_MORE_DOCS; parentDoc = parents
                .nextDoc()) {
            int expectedChildDocs = leaf.reader().document(parentDoc).getField("num_child_docs").numericValue()
                    .intValue();
            hitContext.reset(null, leaf, parentDoc, reader);
            NestedChildrenFilter nestedChildrenFilter = new NestedChildrenFilter(parentFilter, childFilter,
                    hitContext);
            TotalHitCountCollector totalHitCountCollector = new TotalHitCountCollector();
            searcher.search(new ConstantScoreQuery(nestedChildrenFilter), totalHitCountCollector);
            assertThat(totalHitCountCollector.getTotalHits(), equalTo(expectedChildDocs));
            checkedParents++;
        }
    }
    assertThat(checkedParents, equalTo(numParentDocs));
    reader.close();
    dir.close();
}

From source file:org.elasticsearch.search.fetch.innerhits.NestedChildrenFilterTests.java

License:Apache License

@Test
public void testNestedChildrenFilter() throws Exception {
    int numParentDocs = scaledRandomIntBetween(0, 32);
    int maxChildDocsPerParent = scaledRandomIntBetween(8, 16);

    Directory dir = newDirectory();//  ww w .  j  a v  a  2s  .c o m
    RandomIndexWriter writer = new RandomIndexWriter(random(), dir);
    for (int i = 0; i < numParentDocs; i++) {
        int numChildDocs = scaledRandomIntBetween(0, maxChildDocsPerParent);
        List<Document> docs = new ArrayList<>(numChildDocs + 1);
        for (int j = 0; j < numChildDocs; j++) {
            Document childDoc = new Document();
            childDoc.add(new StringField("type", "child", Field.Store.NO));
            docs.add(childDoc);
        }

        Document parenDoc = new Document();
        parenDoc.add(new StringField("type", "parent", Field.Store.NO));
        parenDoc.add(new IntField("num_child_docs", numChildDocs, Field.Store.YES));
        docs.add(parenDoc);
        writer.addDocuments(docs);
    }

    IndexReader reader = writer.getReader();
    writer.close();

    IndexSearcher searcher = new IndexSearcher(reader);
    FetchSubPhase.HitContext hitContext = new FetchSubPhase.HitContext();
    BitSetProducer parentFilter = new QueryBitSetProducer(new TermQuery(new Term("type", "parent")));
    Query childFilter = new TermQuery(new Term("type", "child"));
    int checkedParents = 0;
    final Weight parentsWeight = searcher.createNormalizedWeight(new TermQuery(new Term("type", "parent")),
            false);
    for (LeafReaderContext leaf : reader.leaves()) {
        DocIdSetIterator parents = parentsWeight.scorer(leaf).iterator();
        for (int parentDoc = parents.nextDoc(); parentDoc != DocIdSetIterator.NO_MORE_DOCS; parentDoc = parents
                .nextDoc()) {
            int expectedChildDocs = leaf.reader().document(parentDoc).getField("num_child_docs").numericValue()
                    .intValue();
            hitContext.reset(null, leaf, parentDoc, searcher);
            NestedChildrenQuery nestedChildrenFilter = new NestedChildrenQuery(parentFilter, childFilter,
                    hitContext);
            TotalHitCountCollector totalHitCountCollector = new TotalHitCountCollector();
            searcher.search(new ConstantScoreQuery(nestedChildrenFilter), totalHitCountCollector);
            assertThat(totalHitCountCollector.getTotalHits(), equalTo(expectedChildDocs));
            checkedParents++;
        }
    }
    assertThat(checkedParents, equalTo(numParentDocs));
    reader.close();
    dir.close();
}

From source file:org.elasticsearch.search.fetch.subphase.MatchedQueriesFetchSubPhase.java

License:Apache License

@Override
public void hitsExecute(SearchContext context, InternalSearchHit[] hits) {
    if (hits.length == 0 ||
    // in case the request has only suggest, parsed query is null
            context.parsedQuery() == null) {
        return;//from  ww  w  .  j a  v  a  2s .c o  m
    }
    hits = hits.clone(); // don't modify the incoming hits
    Arrays.sort(hits, (a, b) -> Integer.compare(a.docId(), b.docId()));
    @SuppressWarnings("unchecked")
    List<String>[] matchedQueries = new List[hits.length];
    for (int i = 0; i < matchedQueries.length; ++i) {
        matchedQueries[i] = new ArrayList<>();
    }

    Map<String, Query> namedQueries = new HashMap<>(context.parsedQuery().namedFilters());
    if (context.parsedPostFilter() != null) {
        namedQueries.putAll(context.parsedPostFilter().namedFilters());
    }

    try {
        for (Map.Entry<String, Query> entry : namedQueries.entrySet()) {
            String name = entry.getKey();
            Query query = entry.getValue();
            int readerIndex = -1;
            int docBase = -1;
            Weight weight = context.searcher().createNormalizedWeight(query, false);
            Bits matchingDocs = null;
            final IndexReader indexReader = context.searcher().getIndexReader();
            for (int i = 0; i < hits.length; ++i) {
                InternalSearchHit hit = hits[i];
                int hitReaderIndex = ReaderUtil.subIndex(hit.docId(), indexReader.leaves());
                if (readerIndex != hitReaderIndex) {
                    readerIndex = hitReaderIndex;
                    LeafReaderContext ctx = indexReader.leaves().get(readerIndex);
                    docBase = ctx.docBase;
                    // scorers can be costly to create, so reuse them across docs of the same segment
                    Scorer scorer = weight.scorer(ctx);
                    matchingDocs = Lucene.asSequentialAccessBits(ctx.reader().maxDoc(), scorer);
                }
                if (matchingDocs.get(hit.docId() - docBase)) {
                    matchedQueries[i].add(name);
                }
            }
        }
        for (int i = 0; i < hits.length; ++i) {
            hits[i].matchedQueries(matchedQueries[i].toArray(new String[matchedQueries[i].size()]));
        }
    } catch (IOException e) {
        throw ExceptionsHelper.convertToElastic(e);
    } finally {
        SearchContext.current().clearReleasables(Lifetime.COLLECTION);
    }
}

From source file:org.elasticsearch.search.fetch.subphase.NestedChildrenFilterTests.java

License:Apache License

public void testNestedChildrenFilter() throws Exception {
    int numParentDocs = scaledRandomIntBetween(0, 32);
    int maxChildDocsPerParent = scaledRandomIntBetween(8, 16);

    Directory dir = newDirectory();/*from ww  w .  j a v  a  2s  .c o m*/
    RandomIndexWriter writer = new RandomIndexWriter(random(), dir);
    for (int i = 0; i < numParentDocs; i++) {
        int numChildDocs = scaledRandomIntBetween(0, maxChildDocsPerParent);
        List<Document> docs = new ArrayList<>(numChildDocs + 1);
        for (int j = 0; j < numChildDocs; j++) {
            Document childDoc = new Document();
            childDoc.add(new StringField("type", "child", Field.Store.NO));
            docs.add(childDoc);
        }

        Document parenDoc = new Document();
        parenDoc.add(new StringField("type", "parent", Field.Store.NO));
        parenDoc.add(new LegacyIntField("num_child_docs", numChildDocs, Field.Store.YES));
        docs.add(parenDoc);
        writer.addDocuments(docs);
    }

    IndexReader reader = writer.getReader();
    writer.close();

    IndexSearcher searcher = new IndexSearcher(reader);
    FetchSubPhase.HitContext hitContext = new FetchSubPhase.HitContext();
    BitSetProducer parentFilter = new QueryBitSetProducer(new TermQuery(new Term("type", "parent")));
    Query childFilter = new TermQuery(new Term("type", "child"));
    int checkedParents = 0;
    final Weight parentsWeight = searcher.createNormalizedWeight(new TermQuery(new Term("type", "parent")),
            false);
    for (LeafReaderContext leaf : reader.leaves()) {
        DocIdSetIterator parents = parentsWeight.scorer(leaf).iterator();
        for (int parentDoc = parents.nextDoc(); parentDoc != DocIdSetIterator.NO_MORE_DOCS; parentDoc = parents
                .nextDoc()) {
            int expectedChildDocs = leaf.reader().document(parentDoc).getField("num_child_docs").numericValue()
                    .intValue();
            hitContext.reset(null, leaf, parentDoc, searcher);
            NestedChildrenQuery nestedChildrenFilter = new NestedChildrenQuery(parentFilter, childFilter,
                    hitContext);
            TotalHitCountCollector totalHitCountCollector = new TotalHitCountCollector();
            searcher.search(new ConstantScoreQuery(nestedChildrenFilter), totalHitCountCollector);
            assertThat(totalHitCountCollector.getTotalHits(), equalTo(expectedChildDocs));
            checkedParents++;
        }
    }
    assertThat(checkedParents, equalTo(numParentDocs));
    reader.close();
    dir.close();
}

From source file:org.elasticsearch.search.fetch.subphase.ParentFieldSubFetchPhaseTests.java

License:Apache License

public void testGetParentId() throws Exception {
    ParentFieldMapper fieldMapper = createParentFieldMapper();
    Directory directory = newDirectory();
    IndexWriter indexWriter = new IndexWriter(directory, newIndexWriterConfig());
    Document document = new Document();
    document.add(new SortedDocValuesField(fieldMapper.fieldType().name(), new BytesRef("1")));
    indexWriter.addDocument(document);/*from   w  w  w .  jav  a 2 s  . c  om*/
    indexWriter.close();

    IndexReader indexReader = DirectoryReader.open(directory);
    String id = ParentFieldSubFetchPhase.getParentId(fieldMapper, indexReader.leaves().get(0).reader(), 0);
    assertEquals("1", id);

    indexReader.close();
    directory.close();
}

From source file:org.elasticsearch.search.fetch.subphase.ParentFieldSubFetchPhaseTests.java

License:Apache License

public void testGetParentIdNoParentField() throws Exception {
    ParentFieldMapper fieldMapper = createParentFieldMapper();
    Directory directory = newDirectory();
    IndexWriter indexWriter = new IndexWriter(directory, newIndexWriterConfig());
    Document document = new Document();
    document.add(new SortedDocValuesField("different_field", new BytesRef("1")));
    indexWriter.addDocument(document);//from w  w w  .  j av  a 2s .c  o  m
    indexWriter.close();

    IndexReader indexReader = DirectoryReader.open(directory);
    String id = ParentFieldSubFetchPhase.getParentId(fieldMapper, indexReader.leaves().get(0).reader(), 0);
    assertNull(id);

    indexReader.close();
    directory.close();
}