org.apache.lucene.search.Collector Java Examples

@Override
public List<Collector> create() throws IOException {
  if (firstPhaseGroups.isEmpty()) {
    return Collections.emptyList();
  }

  final List<Collector> collectors = new ArrayList<>(1);
  final FieldType fieldType = field.getType();
  if (fieldType.getNumberType() != null) {
    ValueSource vs = fieldType.getValueSource(field, null);
    Collection<SearchGroup<MutableValue>> v = GroupConverter.toMutable(field, firstPhaseGroups);
    secondPassCollector = new TopGroupsCollector<>(new ValueSourceGroupSelector(vs, new HashMap<>()),
        v, groupSort, withinGroupSort, maxDocPerGroup, needMaxScore
    );
  } else {
    secondPassCollector = new TopGroupsCollector<>(new TermGroupSelector(field.getName()),
        firstPhaseGroups, groupSort, withinGroupSort, maxDocPerGroup, needMaxScore
    );
  }
  collectors.add(secondPassCollector);
  return collectors;
}
 

protected Collector getSecureCollector(final Collector collector) {
  return new Collector() {

    @Override
    public void setScorer(Scorer scorer) throws IOException {
      collector.setScorer(scorer);
    }

    @Override
    public void setNextReader(AtomicReaderContext context) throws IOException {
      Object key = context.reader().getCoreCacheKey();
      AtomicReaderContext atomicReaderContext = _leaveMap.get(key);
      collector.setNextReader(atomicReaderContext);
    }

    @Override
    public void collect(int doc) throws IOException {
      collector.collect(doc);
    }

    @Override
    public boolean acceptsDocsOutOfOrder() {
      return collector.acceptsDocsOutOfOrder();
    }
  };
}
 

private Collector getInsanityWrapper(final String field, Collector collector) {
  SchemaField sf = searcher.getSchema().getFieldOrNull(field);
  if (sf != null && !sf.hasDocValues() && !sf.multiValued() && sf.getType().getNumberType() != null) {
    // it's a single-valued numeric field: we must currently create insanity :(
    // there isn't a GroupedFacetCollector that works on numerics right now...
    return new FilterCollector(collector) {
      @Override
      public LeafCollector getLeafCollector(LeafReaderContext context) throws IOException {
        LeafReader insane = Insanity.wrapInsanity(context.reader(), field);
        return in.getLeafCollector(insane.getContext());
      }
    };
  } else {
    return collector;
  }
}
 

/**
 * Invokes search with the specified filter and collector.  
 * If a time limit has been specified, wrap the collector in a TimeLimitingCollector
 */
private void searchWithTimeLimiter(final Filter luceneFilter, Collector collector) throws IOException {
  if (cmd.getTimeAllowed() > 0) {
    if (timeLimitingCollector == null) {
      timeLimitingCollector = new TimeLimitingCollector(collector, TimeLimitingCollector.getGlobalCounter(), cmd.getTimeAllowed());
    } else {
      /*
       * This is so the same timer can be used for grouping's multiple phases.   
       * We don't want to create a new TimeLimitingCollector for each phase because that would 
       * reset the timer for each phase.  If time runs out during the first phase, the 
       * second phase should timeout quickly.
       */
      timeLimitingCollector.setCollector(collector);
    }
    collector = timeLimitingCollector;
  }
  try {
    searcher.search(QueryUtils.combineQueryAndFilter(query, luceneFilter), collector);
  } catch (TimeLimitingCollector.TimeExceededException | ExitableDirectoryReader.ExitingReaderException x) {
    log.warn("Query: {}; {}", query, x.getMessage());
    qr.setPartialResults(true);
  }
}
 

@Override
protected Collector createFirstPassCollector() throws IOException {
  DocSet groupFilt = searcher.getDocSet(query);
  int groupDocsToCollect = getMax(groupOffset, docsPerGroup, maxDoc);
  Collector subCollector;
  if (withinGroupSort == null || withinGroupSort.equals(Sort.RELEVANCE)) {
    subCollector = topCollector = TopScoreDocCollector.create(groupDocsToCollect, Integer.MAX_VALUE);
  } else {
    topCollector = TopFieldCollector.create(searcher.weightSort(withinGroupSort), groupDocsToCollect, Integer.MAX_VALUE);
    if (needScores) {
      maxScoreCollector = new MaxScoreCollector();
      subCollector = MultiCollector.wrap(topCollector, maxScoreCollector);
    } else {
      subCollector = topCollector;
    }
  }
  collector = new FilterCollector(groupFilt, subCollector);
  return collector;
}
 

private DocSet computeGroupedDocSet(Query query, ProcessedFilter filter, List<Collector> collectors) throws IOException {
  @SuppressWarnings({"rawtypes"})
  Command firstCommand = commands.get(0);
  String field = firstCommand.getKey();
  SchemaField sf = searcher.getSchema().getField(field);
  FieldType fieldType = sf.getType();
  
  @SuppressWarnings({"rawtypes"})
  final AllGroupHeadsCollector allGroupHeadsCollector;
  if (fieldType.getNumberType() != null) {
    ValueSource vs = fieldType.getValueSource(sf, null);
    allGroupHeadsCollector = AllGroupHeadsCollector.newCollector(new ValueSourceGroupSelector(vs, new HashMap<>()),
        firstCommand.getWithinGroupSort());
  } else {
    allGroupHeadsCollector
        = AllGroupHeadsCollector.newCollector(new TermGroupSelector(firstCommand.getKey()), firstCommand.getWithinGroupSort());
  }
  if (collectors.isEmpty()) {
    searchWithTimeLimiter(query, filter, allGroupHeadsCollector);
  } else {
    collectors.add(allGroupHeadsCollector);
    searchWithTimeLimiter(query, filter, MultiCollector.wrap(collectors.toArray(new Collector[collectors.size()])));
  }

  return new BitDocSet(allGroupHeadsCollector.retrieveGroupHeads(searcher.maxDoc()));
}
 

/**
 * Creates a human-friendly representation of the Collector name.
 *
 * Bucket Collectors use the aggregation name in their toString() method,
 * which makes the profiled output a bit nicer.
 *
 * @param c The Collector to derive a name from
 * @return  A (hopefully) prettier name
 */
private String deriveCollectorName(Collector c) {
    String s = c.getClass().getSimpleName();

    // MutiCollector which wraps multiple BucketCollectors is generated
    // via an anonymous class, so this corrects the lack of a name by
    // asking the enclosingClass
    if (s.equals("")) {
        s = c.getClass().getEnclosingClass().getSimpleName();
    }

    // Aggregation collector toString()'s include the user-defined agg name
    if (reason.equals(CollectorResult.REASON_AGGREGATION) || reason.equals(CollectorResult.REASON_AGGREGATION_GLOBAL)) {
        s += ": [" + c.toString() + "]";
    }
    return s;
}
 

/**
 * Get docIds from the text inverted index for a given raw value
 * @param value value to look for in the inverted index
 * @return docIDs in bitmap
 */
@Override
public MutableRoaringBitmap getDocIds(Object value) {
  String searchQuery = (String) value;
  MutableRoaringBitmap docIds = new MutableRoaringBitmap();
  Collector docIDCollector = new LuceneDocIdCollector(docIds, _docIdTranslator);
  try {
    // Lucene Query Parser is JavaCC based. It is stateful and should
    // be instantiated per query. Analyzer on the other hand is stateless
    // and can be created upfront.
    QueryParser parser = new QueryParser(_column, _standardAnalyzer);
    Query query = parser.parse(searchQuery);
    _indexSearcher.search(query, docIDCollector);
    return docIds;
  } catch (Exception e) {
    String msg = "Caught excepttion while searching the text index for column:" + _column + " search query:" + searchQuery;
    throw new RuntimeException(msg, e);
  }
}
 

@Override
public List<Collector> create() throws IOException {
  Collector subCollector;
  if (sort == null || sort.equals(Sort.RELEVANCE)) {
    subCollector = topDocsCollector = TopScoreDocCollector.create(docsToCollect, Integer.MAX_VALUE);
  } else {
    topDocsCollector = TopFieldCollector.create(sort, docsToCollect, Integer.MAX_VALUE);
    if (needScores) {
      maxScoreCollector = new MaxScoreCollector();
      subCollector = MultiCollector.wrap(topDocsCollector, maxScoreCollector);
    } else {
      subCollector = topDocsCollector;
    }
  }
  filterCollector = new FilterCollector(docSet, subCollector);
  return Arrays.asList((Collector) filterCollector);
}
 

/** Utility method, to search and also collect all hits
 *  into the provided {@link Collector}. */
public static TopDocs searchAfter(IndexSearcher searcher, ScoreDoc after, Query q, int n, Sort sort, boolean doDocScores, Collector fc) throws IOException {
  if (sort == null) {
    throw new IllegalArgumentException("sort must not be null");
  }
  return doSearch(searcher, after, q, n, sort, doDocScores, fc);
}
 

private KeyIterable<ShardId, Row> doSearch(TopFieldCollector topFieldCollector,
                                           Float minScore,
                                           Query query) throws IOException {
    Collector collector = topFieldCollector;
    if (minScore != null) {
        collector = new MinimumScoreCollector(collector, minScore);
    }
    collector = new KillableCollector(collector, this::raiseIfKilled);
    searcher.search(query, collector);
    ScoreDoc[] scoreDocs = topFieldCollector.topDocs().scoreDocs;
    if (doDocsScores) {
        TopFieldCollector.populateScores(scoreDocs, searcher, query);
    }
    return scoreDocToIterable(scoreDocs);
}
 

@Override
public void preProcess(SearchContext context) {
    if (context.aggregations() != null) {
        AggregationContext aggregationContext = new AggregationContext(context);
        context.aggregations().aggregationContext(aggregationContext);

        List<Aggregator> collectors = new ArrayList<>();
        Aggregator[] aggregators;
        try {
            AggregatorFactories factories = context.aggregations().factories();
            aggregators = factories.createTopLevelAggregators(aggregationContext);
            for (int i = 0; i < aggregators.length; i++) {
                if (aggregators[i] instanceof GlobalAggregator == false) {
                    collectors.add(aggregators[i]);
                }
            }
            context.aggregations().aggregators(aggregators);
            if (!collectors.isEmpty()) {
                Collector collector = BucketCollector.wrap(collectors);
                ((BucketCollector)collector).preCollection();
                if (context.getProfilers() != null) {
                    // TODO: report on child aggs as well
                    List<InternalProfileCollector> emptyList = Collections.emptyList();
                    collector = new InternalProfileCollector(collector, CollectorResult.REASON_AGGREGATION, emptyList);
                }
                context.queryCollectors().put(AggregationPhase.class, collector);
            }
        } catch (IOException e) {
            throw new AggregationInitializationException("Could not initialize aggregators", e);
        }
    }
}
 

@Override
public Collector newCollector() throws IOException {
  TopFieldCollector collector = TopFieldCollector.create(_sort, _numHitsToCollect, _after, true, true, false, true);
  Collector col = new StopExecutionCollector(collector, _running);
  if (_runSlow) {
    return new SlowCollector(col);
  }
  return col;
}
 

private TopScoreDocCollector getTopScoreDocCollector(Collector collector) {
  if (collector instanceof SlowCollector) {
    SlowCollector slowCollector = (SlowCollector) collector;
    return getTopScoreDocCollector(slowCollector.getCollector());
  } else if (collector instanceof StopExecutionCollector) {
    StopExecutionCollector stopExecutionCollector = (StopExecutionCollector) collector;
    return getTopScoreDocCollector(stopExecutionCollector.getCollector());
  } else if (collector instanceof TopScoreDocCollector) {
    TopScoreDocCollector topScoreDocCollector = (TopScoreDocCollector) collector;
    return topScoreDocCollector;
  } else {
    throw new RuntimeException("Collector type [" + collector + "] not supported.");
  }
}
 

@Override
public Collector newCollector() {
  TopScoreDocCollector collector = TopScoreDocCollector.create(_numHitsToCollect, _after, true);
  Collector col = new StopExecutionCollector(collector, _running);
  if (_runSlow) {
    return new SlowCollector(col);
  }
  return col;
}
 

public LindenDocsCollector(Collector collector) {
  if (!(collector instanceof TopDocsCollector) && !(collector instanceof EarlyTerminationCollector)) {
    throw new RuntimeException("Unsupported collector class in LindenDocsCollector: " + collector.getClass().getName());
  }
  hitCollector = collector;
  wrappedCollector = collector;
}
 

private TopFieldCollector getTopFieldCollector(Collector collector) {
  if (collector instanceof SlowCollector) {
    SlowCollector slowCollector = (SlowCollector) collector;
    return getTopFieldCollector(slowCollector.getCollector());
  } else if (collector instanceof StopExecutionCollector) {
    StopExecutionCollector stopExecutionCollector = (StopExecutionCollector) collector;
    return getTopFieldCollector(stopExecutionCollector.getCollector());
  } else if (collector instanceof TopFieldCollector) {
    TopFieldCollector topFieldCollector = (TopFieldCollector) collector;
    return topFieldCollector;
  } else {
    throw new RuntimeException("Collector type [" + collector + "] not supported.");
  }
}
 

/** Utility method, to search and also collect all hits
 *  into the provided {@link Collector}. */
public static TopFieldDocs search(IndexSearcher searcher, Query q, int n, Sort sort, boolean doDocScores, Collector fc) throws IOException {
  if (sort == null) {
    throw new IllegalArgumentException("sort must not be null");
  }
  return (TopFieldDocs) doSearch(searcher, null, q, n, sort, doDocScores, fc);
}
 

private Callable<Void> newSearchCallable(final Weight weight, final Collector collector, final AtomicReaderContext ctx) {
  return new Callable<Void>() {
    @Override
    public Void call() throws Exception {
      runSearch(weight, collector, ctx);
      return null;
    }
  };
}
 

/** Utility method, to search and also collect all hits
 *  into the provided {@link Collector}. */
public static TopDocs searchAfter(IndexSearcher searcher, ScoreDoc after, Query q, int n, Sort sort, Collector fc) throws IOException {
  if (sort == null) {
    throw new IllegalArgumentException("sort must not be null");
  }
  return doSearch(searcher, after, q, n, sort, false, fc);
}
 

/** Utility method, to search and also collect all hits
 *  into the provided {@link Collector}. */
public static TopFieldDocs search(IndexSearcher searcher, Query q, int n, Sort sort, Collector fc) throws IOException {
  if (sort == null) {
    throw new IllegalArgumentException("sort must not be null");
  }
  return (TopFieldDocs) doSearch(searcher, null, q, n, sort, false, fc);
}
 

private DocSet computeDocSet(Query query, ProcessedFilter filter, List<Collector> collectors) throws IOException {
  int maxDoc = searcher.maxDoc();
  final DocSetCollector docSetCollector = new DocSetCollector(maxDoc);
  List<Collector> allCollectors = new ArrayList<>(collectors);
  allCollectors.add(docSetCollector);
  searchWithTimeLimiter(query, filter, MultiCollector.wrap(allCollectors));
  return DocSetUtil.getDocSet( docSetCollector, searcher );
}
 

DrillSidewaysQuery(Query baseQuery, Collector drillDownCollector, Collector[] drillSidewaysCollectors, Query[] drillDownQueries, boolean scoreSubDocsAtOnce) {
  this.baseQuery = Objects.requireNonNull(baseQuery);
  this.drillDownCollector = drillDownCollector;
  this.drillSidewaysCollectors = drillSidewaysCollectors;
  this.drillDownQueries = drillDownQueries;
  this.scoreSubDocsAtOnce = scoreSubDocsAtOnce;
}
 

private static LongBitSet findFieldOrdinalsMatchingQuery(Query q, String field, SolrIndexSearcher searcher, SortedSetDocValues docValues) throws IOException {
  final LongBitSet fromOrdBitSet = new LongBitSet(docValues.getValueCount());
  final Collector fromCollector = new MultiValueTermOrdinalCollector(field, docValues, fromOrdBitSet);

  searcher.search(q, fromCollector);

  return fromOrdBitSet;
}
 

@Override
default ScoreMode scoreMode() {
  final LongObjectMap<Collector> groups = getGroups();
  if (groups.isEmpty()) {
    return ScoreMode.COMPLETE; // doesn't matter?
  } else {
    return groups.iterator().next().value.scoreMode(); // we assume all the collectors should have the same nature
  }
}
 

@Override
protected Collector createFirstPassCollector() throws IOException {
  // Ok we don't want groups, but do want a total count
  if (actualGroupsToFind <= 0) {
    fallBackCollector = new TotalHitCountCollector();
    return fallBackCollector;
  }

  groupSort = groupSort == null ? Sort.RELEVANCE : groupSort;
  firstPass = new FirstPassGroupingCollector<>(new TermGroupSelector(groupBy), groupSort, actualGroupsToFind);
  return firstPass;
}
 

public static void collectSortedDocSet(DocSet docs, IndexReader reader, Collector collector) throws IOException {
  // TODO add SortedDocSet sub-interface and take that.
  // TODO collectUnsortedDocSet: iterate segment, then all docSet per segment.

  final List<LeafReaderContext> leaves = reader.leaves();
  final Iterator<LeafReaderContext> ctxIt = leaves.iterator();
  int segBase = 0;
  int segMax;
  int adjustedMax = 0;
  LeafReaderContext ctx = null;
  LeafCollector leafCollector = null;
  for (DocIterator docsIt = docs.iterator(); docsIt.hasNext(); ) {
    final int doc = docsIt.nextDoc();
    if (doc >= adjustedMax) {
      do {
        ctx = ctxIt.next();
        segBase = ctx.docBase;
        segMax = ctx.reader().maxDoc();
        adjustedMax = segBase + segMax;
      } while (doc >= adjustedMax);
      leafCollector = collector.getLeafCollector(ctx);
    }
    if (doc < segBase) {
      throw new IllegalStateException("algorithm expects sorted DocSet but wasn't: " + docs.getClass());
    }
    leafCollector.collect(doc - segBase);  // per-seg collectors
  }
}
 

@Override
protected Collector createFirstPassCollector() throws IOException {
  // Ok we don't want groups, but do want a total count
  if (actualGroupsToFind <= 0) {
    fallBackCollector = new TotalHitCountCollector();
    return fallBackCollector;
  }

  groupSort = groupSort == null ? Sort.RELEVANCE : groupSort;
  firstPass = new FirstPassGroupingCollector<>(newSelector(), searcher.weightSort(groupSort), actualGroupsToFind);
  return firstPass;
}
 

/** Sets the last delegate in a chain of DelegatingCollectors */
public void setLastDelegate(Collector delegate) {
  DelegatingCollector ptr = this;
  for(; ptr.getDelegate() instanceof DelegatingCollector; ptr = (DelegatingCollector)ptr.getDelegate());
  ptr.setDelegate(delegate);
  setLastDelegateCount++;
}
 

/**
 * Invokes search with the specified filter and collector.  
 * If a time limit has been specified then wrap the collector in the TimeLimitingCollector
 */
private void searchWithTimeLimiter(Query query, 
                                   ProcessedFilter filter, 
                                   Collector collector) throws IOException {
  if (queryCommand.getTimeAllowed() > 0 ) {
    collector = new TimeLimitingCollector(collector, TimeLimitingCollector.getGlobalCounter(), queryCommand.getTimeAllowed());
  }

  TotalHitCountCollector hitCountCollector = new TotalHitCountCollector();
  if (includeHitCount) {
    collector = MultiCollector.wrap(collector, hitCountCollector);
  }

  query = QueryUtils.combineQueryAndFilter(query, filter.filter);

  if (filter.postFilter != null) {
    filter.postFilter.setLastDelegate(collector);
    collector = filter.postFilter;
  }

  try {
    searcher.search(query, collector);
  } catch (TimeLimitingCollector.TimeExceededException | ExitableDirectoryReader.ExitingReaderException x) {
    partialResults = true;
    log.warn("Query: {}; {}", query, x.getMessage());
  }

  if (includeHitCount) {
    totalHitCount = hitCountCollector.getTotalHits();
  }
}