Java Code Examples for org.apache.lucene.index.ReaderUtil#subIndex()

@Nullable
@Override
public Row apply(@Nullable ScoreDoc input) {
    if (input == null) {
        return null;
    }
    FieldDoc fieldDoc = (FieldDoc) input;
    scorer.score(fieldDoc.score);
    for (OrderByCollectorExpression orderByCollectorExpression : orderByCollectorExpressions) {
        orderByCollectorExpression.setNextFieldDoc(fieldDoc);
    }
    List<LeafReaderContext> leaves = indexReader.leaves();
    int readerIndex = ReaderUtil.subIndex(fieldDoc.doc, leaves);
    LeafReaderContext subReaderContext = leaves.get(readerIndex);
    int subDoc = fieldDoc.doc - subReaderContext.docBase;
    for (LuceneCollectorExpression<?> expression : expressions) {
        expression.setNextReader(subReaderContext);
        expression.setNextDocId(subDoc);
    }
    return inputRow;
}
 

@Override
public Object[] apply(Integer docId) {
    List<LeafReaderContext> leaves = searcher.reader().leaves();
    int readerIndex = ReaderUtil.subIndex(docId, leaves);
    LeafReaderContext leafContext = leaves.get(readerIndex);
    int subDoc = docId - leafContext.docBase;
    for (LuceneCollectorExpression<?> expression : expressions) {
        try {
            expression.setNextReader(leafContext);
            expression.setNextDocId(subDoc);
        } catch (IOException e) {
            throw new UncheckedIOException(e);
        }
    }
    Object[] cells = new Object[inputs.size()];
    for (int i = 0; i < cells.length; i++) {
        cells[i] = inputs.get(i).value();
    }
    return cells;
}
 

@Nullable
@Override
public Row apply(@Nullable ScoreDoc input) {
    onScoreDoc.run();
    if (input == null) {
        return null;
    }
    FieldDoc fieldDoc = (FieldDoc) input;
    scorer.score(fieldDoc.score);
    for (int i = 0; i < orderByCollectorExpressions.size(); i++) {
        orderByCollectorExpressions.get(i).setNextFieldDoc(fieldDoc);
    }
    List<LeafReaderContext> leaves = indexReader.leaves();
    int readerIndex = ReaderUtil.subIndex(fieldDoc.doc, leaves);
    LeafReaderContext subReaderContext = leaves.get(readerIndex);
    int subDoc = fieldDoc.doc - subReaderContext.docBase;
    for (LuceneCollectorExpression<?> expression : expressions) {
        try {
            expression.setNextReader(subReaderContext);
            expression.setNextDocId(subDoc);
        } catch (IOException e) {
            throw new UncheckedIOException(e);
        }
    }
    return inputRow;
}
 

Object objectVal(int topDocId) throws IOException {
  // lookup segment level stuff:
  int segIdx = ReaderUtil.subIndex(topDocId, readerContexts);
  LeafReaderContext rcontext = readerContexts.get(segIdx);
  int segDocId = topDocId - rcontext.docBase;
  // unfortunately Lucene 7.0 requires forward only traversal (with no reset method).
  //   So we need to track our last docId (per segment) and re-fetch the FunctionValues. :-(
  FunctionValues functionValues = functionValuesPerSeg[segIdx];
  if (functionValues == null || segDocId < functionValuesDocIdPerSeg[segIdx]) {
    functionValues = functionValuesPerSeg[segIdx] = valueSource.getValues(fContext, rcontext);
  }
  functionValuesDocIdPerSeg[segIdx] = segDocId;

  // get value:
  return functionValues.objectVal(segDocId);
}
 

/** 
 * Returns the weight for the current <code>docId</code> as computed 
 * by the <code>weightsValueSource</code>
 * */
@Override
protected long getWeight(Document doc, int docId) throws IOException {
  if (currentWeightValues == null) {
    return 0;
  }
  int subIndex = ReaderUtil.subIndex(docId, starts);
  if (subIndex != currentLeafIndex) {
    currentLeafIndex = subIndex;
    currentWeightValues = weightsValueSource.getValues(leaves.get(currentLeafIndex), null);
  }
  if (currentWeightValues.advanceExact(docId - starts[subIndex]))
    return currentWeightValues.longValue();
  else
    return 0;

}
 

@SuppressWarnings({"unchecked"})
Object objectVal(int topDocId) throws IOException {
  // lookup segment level stuff:
  int segIdx = ReaderUtil.subIndex(topDocId, readerContexts);
  LeafReaderContext rcontext = readerContexts.get(segIdx);
  int segDocId = topDocId - rcontext.docBase;
  // unfortunately Lucene 7.0 requires forward only traversal (with no reset method).
  //   So we need to track our last docId (per segment) and re-fetch the FunctionValues. :-(
  FunctionValues functionValues = functionValuesPerSeg[segIdx];
  if (functionValues == null || segDocId < functionValuesDocIdPerSeg[segIdx]) {
    functionValues = functionValuesPerSeg[segIdx] = valueSource.getValues(fContext, rcontext);
  }
  functionValuesDocIdPerSeg[segIdx] = segDocId;

  // get value:
  return functionValues.objectVal(segDocId);
}
 

@Override
public void transform(SolrDocument doc, int docid) {
  // This is only good for random-access functions

  try {

    // TODO: calculate this stuff just once across diff functions
    int idx = ReaderUtil.subIndex(docid, readerContexts);
    LeafReaderContext rcontext = readerContexts.get(idx);
    @SuppressWarnings({"unchecked"})
    FunctionValues values = valueSource.getValues(fcontext, rcontext);
    int localId = docid - rcontext.docBase;
    setValue(doc,values.objectVal(localId));
  } catch (IOException e) {
    throw new SolrException(SolrException.ErrorCode.SERVER_ERROR, "exception at docid " + docid + " for valuesource " + valueSource, e);
  }
}
 

public static LTRScoringQuery.FeatureInfo[] extractFeaturesInfo(LTRScoringQuery.ModelWeight modelWeight,
    int docid,
    Float originalDocScore,
    List<LeafReaderContext> leafContexts)
        throws IOException {
  final int n = ReaderUtil.subIndex(docid, leafContexts);
  final LeafReaderContext atomicContext = leafContexts.get(n);
  final int deBasedDoc = docid - atomicContext.docBase;
  final LTRScoringQuery.ModelWeight.ModelScorer r = modelWeight.scorer(atomicContext);
  if ( (r == null) || (r.iterator().advance(deBasedDoc) != deBasedDoc) ) {
    return new LTRScoringQuery.FeatureInfo[0];
  } else {
    if (originalDocScore != null) {
      // If results have not been reranked, the score passed in is the original query's
      // score, which some features can use instead of recalculating it
      r.getDocInfo().setOriginalDocScore(originalDocScore);
    }
    r.score();
    return modelWeight.getFeaturesInfo();
  }
}
 

private LTRScoringQuery.ModelWeight performQuery(TopDocs hits,
    IndexSearcher searcher, int docid, LTRScoringQuery model) throws IOException,
    ModelException {
  final List<LeafReaderContext> leafContexts = searcher.getTopReaderContext()
      .leaves();
  final int n = ReaderUtil.subIndex(hits.scoreDocs[0].doc, leafContexts);
  final LeafReaderContext context = leafContexts.get(n);
  final int deBasedDoc = hits.scoreDocs[0].doc - context.docBase;

  final Weight weight = searcher.createWeight(searcher.rewrite(model), ScoreMode.COMPLETE, 1);
  final Scorer scorer = weight.scorer(context);

  // rerank using the field final-score
  scorer.iterator().advance(deBasedDoc);
  scorer.score();
  assertTrue(weight instanceof LTRScoringQuery.ModelWeight);
  final LTRScoringQuery.ModelWeight modelWeight = (LTRScoringQuery.ModelWeight) weight;
  return modelWeight;

}
 

private LTRScoringQuery.ModelWeight performQuery(TopDocs hits,
    IndexSearcher searcher, int docid, LTRScoringQuery model) throws IOException,
    ModelException {
  final List<LeafReaderContext> leafContexts = searcher.getTopReaderContext()
      .leaves();
  final int n = ReaderUtil.subIndex(hits.scoreDocs[0].doc, leafContexts);
  final LeafReaderContext context = leafContexts.get(n);
  final int deBasedDoc = hits.scoreDocs[0].doc - context.docBase;

  final Weight weight = searcher.createWeight(searcher.rewrite(model), ScoreMode.COMPLETE, 1);
  final Scorer scorer = weight.scorer(context);

  // rerank using the field final-score
  scorer.iterator().advance(deBasedDoc);
  scorer.score();

  // assertEquals(42.0f, score, 0.0001);
  // assertTrue(weight instanceof AssertingWeight);
  // (AssertingIndexSearcher)
  assertTrue(weight instanceof LTRScoringQuery.ModelWeight);
  final LTRScoringQuery.ModelWeight modelWeight = (LTRScoringQuery.ModelWeight) weight;
  return modelWeight;

}
 

public StreamBucket collect(IntContainer docIds) {
    StreamBucket.Builder builder = new StreamBucket.Builder(streamers, ramAccounting);
    for (IntCursor cursor : docIds) {
        int docId = cursor.value;
        int readerIndex = ReaderUtil.subIndex(docId, readerContexts);
        LeafReaderContext subReaderContext = readerContexts.get(readerIndex);
        try {
            setNextDocId(subReaderContext, docId - subReaderContext.docBase);
        } catch (IOException e) {
            Exceptions.rethrowRuntimeException(e);
        }
        builder.add(row);
    }
    return builder.build();
}
 

@SuppressWarnings({"unchecked"})
private String getStringVal(SolrQueryRequest sqr, String field, int doc) throws IOException {
  SchemaField sf = sqr.getSchema().getField(field);
  ValueSource vs = sf.getType().getValueSource(sf, null);
  @SuppressWarnings({"rawtypes"})
  Map context = ValueSource.newContext(sqr.getSearcher());
  vs.createWeight(context, sqr.getSearcher());
  IndexReaderContext topReaderContext = sqr.getSearcher().getTopReaderContext();
  List<LeafReaderContext> leaves = topReaderContext.leaves();
  int idx = ReaderUtil.subIndex(doc, leaves);
  LeafReaderContext leaf = leaves.get(idx);
  FunctionValues vals = vs.getValues(context, leaf);
  return vals.strVal(doc-leaf.docBase);
}
 

/**
 * This will fetch and add the docValues fields to a given SolrDocument/SolrInputDocument
 *
 * @param doc
 *          A SolrDocument or SolrInputDocument instance where docValues will be added
 * @param docid
 *          The lucene docid of the document to be populated
 * @param fields
 *          The fields with docValues to populate the document with.
 *          DocValues fields which do not exist or not decodable will be ignored.
 */
public void decorateDocValueFields(@SuppressWarnings("rawtypes") SolrDocumentBase doc, int docid, Set<String> fields)
    throws IOException {
  final List<LeafReaderContext> leafContexts = searcher.getLeafContexts();
  final int subIndex = ReaderUtil.subIndex(docid, leafContexts);
  final int localId = docid - leafContexts.get(subIndex).docBase;
  final LeafReader leafReader = leafContexts.get(subIndex).reader();
  for (String fieldName : fields) {
    Object fieldValue = decodeDVField(localId, leafReader, fieldName);
    if (fieldValue != null) {
      doc.setField(fieldName, fieldValue);
    }
  }
}
 

@Override
public Explanation explain(IndexSearcher searcher,
    Explanation firstPassExplanation, int docID) throws IOException {

  final List<LeafReaderContext> leafContexts = searcher.getTopReaderContext()
      .leaves();
  final int n = ReaderUtil.subIndex(docID, leafContexts);
  final LeafReaderContext context = leafContexts.get(n);
  final int deBasedDoc = docID - context.docBase;
  final Weight modelWeight = searcher.createWeight(searcher.rewrite(scoringQuery),
      ScoreMode.COMPLETE, 1);
  return modelWeight.explain(context, deBasedDoc);
}
 

private Document getParentDoc(IndexReader reader, BitSetProducer parents, int childDocID) throws IOException {
  final List<LeafReaderContext> leaves = reader.leaves();
  final int subIndex = ReaderUtil.subIndex(childDocID, leaves);
  final LeafReaderContext leaf = leaves.get(subIndex);
  final BitSet bits = parents.getBitSet(leaf);
  return leaf.reader().document(bits.nextSetBit(childDocID - leaf.docBase));
}
 

@Override
public Explanation explain(IndexSearcher searcher, Explanation firstPassExplanation, int docID) throws IOException {
  Explanation superExpl = super.explain(searcher, firstPassExplanation, docID);

  List<LeafReaderContext> leaves = searcher.getIndexReader().leaves();
  int subReader = ReaderUtil.subIndex(docID, leaves);
  LeafReaderContext readerContext = leaves.get(subReader);
  int docIDInSegment = docID - readerContext.docBase;

  return expression.getDoubleValuesSource(bindings).explain(readerContext, docIDInSegment, superExpl);
}
 

private void assertGaps(IntervalsSource source, int doc, String field, int[] expectedGaps) throws IOException {
  int ord = ReaderUtil.subIndex(doc, searcher.getIndexReader().leaves());
  LeafReaderContext ctx = searcher.getIndexReader().leaves().get(ord);
  IntervalIterator it = source.intervals(field, ctx);
  doc = doc - ctx.docBase;
  assertEquals(doc, it.advance(doc));
  for (int expectedGap : expectedGaps) {
    if (it.nextInterval() == IntervalIterator.NO_MORE_INTERVALS) {
      fail("Unexpected interval " + it);
    }
    assertEquals(expectedGap, it.gaps());
  }
}
 

public void collect(IntContainer docIds, StreamBucket.Builder builder) throws IOException {
    for (IntCursor cursor : docIds) {
        final int docId = cursor.value;
        int readerIndex = ReaderUtil.subIndex(docId, readerContexts);
        LeafReaderContext subReaderContext = readerContexts.get(readerIndex);
        setNextReader(subReaderContext);
        setNextDocId(docId - subReaderContext.docBase);
        builder.add(row);
    }
}
 

private MatchesIterator getMatches(IntervalsSource source, int doc, String field) throws IOException {
  int ord = ReaderUtil.subIndex(doc, searcher.getIndexReader().leaves());
  LeafReaderContext ctx = searcher.getIndexReader().leaves().get(ord);
  return source.matches(field, ctx, doc - ctx.docBase);
}
 

/** Expert: low-level implementation method
 * Returns an Explanation that describes how <code>doc</code> scored against
 * <code>weight</code>.
 *
 * <p>This is intended to be used in developing Similarity implementations,
 * and, for good performance, should not be displayed with every hit.
 * Computing an explanation is as expensive as executing the query over the
 * entire index.
 * <p>Applications should call {@link IndexSearcher#explain(Query, int)}.
 * @throws TooManyClauses If a query would exceed
 *         {@link IndexSearcher#getMaxClauseCount()} clauses.
 */
protected Explanation explain(Weight weight, int doc) throws IOException {
  int n = ReaderUtil.subIndex(doc, leafContexts);
  final LeafReaderContext ctx = leafContexts.get(n);
  int deBasedDoc = doc - ctx.docBase;
  final Bits liveDocs = ctx.reader().getLiveDocs();
  if (liveDocs != null && liveDocs.get(deBasedDoc) == false) {
    return Explanation.noMatch("Document " + doc + " is deleted");
  }
  return weight.explain(ctx, deBasedDoc);
}