org.apache.solr.search.grouping.distributed.command.QueryCommand Maven / Gradle / Ivy

Show more of this group Show more artifacts with this name
Show all versions of solr-core Show documentation
Apache Solr (module: core)
There is a newer version: 9.7.0
/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.solr.search.grouping.distributed.command;

import java.io.IOException;
import java.util.Arrays;
import java.util.List;
import org.apache.lucene.search.Collector;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.MultiCollector;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.Sort;
import org.apache.lucene.search.TopDocs;
import org.apache.lucene.search.TopDocsCollector;
import org.apache.lucene.search.TopFieldCollector;
import org.apache.lucene.search.TopScoreDocCollector;
import org.apache.solr.request.SolrQueryRequest;
import org.apache.solr.search.DocSet;
import org.apache.solr.search.MaxScoreCollector;
import org.apache.solr.search.QParser;
import org.apache.solr.search.SolrIndexSearcher;
import org.apache.solr.search.SyntaxError;
import org.apache.solr.search.grouping.Command;
import org.apache.solr.search.grouping.collector.FilterCollector;

/** */
public class QueryCommand implements Command {

  public static class Builder {

    private Sort sort;
    private String queryString;
    private Query query;
    private Query mainQuery;
    private DocSet docSet;
    private Integer docsToCollect;
    private boolean needScores;

    public Builder setSort(Sort sort) {
      this.sort = sort;
      return this;
    }

    /**
     * Sets the group query.
     *
     * @param query The {@link Query} used for grouping
     * @return this
     */
    public Builder setQuery(Query query) {
      this.query = query;
      return this;
    }

    /**
     * Sets the main query used for fetching results. This is mainly used for computing the scores.
     *
     * @param mainQuery The top-level query
     * @return this
     */
    public Builder setMainQuery(Query mainQuery) {
      this.mainQuery = mainQuery;
      return this;
    }

    /**
     * Sets the group query from the specified groupQueryString. The groupQueryString is parsed into
     * a query.
     *
     * @param groupQueryString The group query string to parse
     * @param request The current request
     * @return this
     */
    public Builder setQuery(String groupQueryString, SolrQueryRequest request) throws SyntaxError {
      QParser parser = QParser.getParser(groupQueryString, request);
      this.queryString = groupQueryString;
      return setQuery(parser.getQuery());
    }

    public Builder setDocSet(DocSet docSet) {
      this.docSet = docSet;
      return this;
    }

    /**
     * Sets the docSet based on the created {@link DocSet}
     *
     * @param searcher The searcher executing the
     * @return this
     * @throws IOException If I/O related errors occur.
     */
    public Builder setDocSet(SolrIndexSearcher searcher) throws IOException {
      return setDocSet(searcher.getDocSet(query));
    }

    public Builder setDocsToCollect(int docsToCollect) {
      this.docsToCollect = docsToCollect;
      return this;
    }

    public Builder setNeedScores(boolean needScores) {
      this.needScores = needScores;
      return this;
    }

    public QueryCommand build() {
      if (sort == null
          || query == null
          || docSet == null
          || docsToCollect == null
          || mainQuery == null) {
        throw new IllegalStateException("All fields must be set");
      }

      return new QueryCommand(
          sort, query, docsToCollect, needScores, docSet, queryString, mainQuery);
    }
  }

  private final Sort sort;
  private final Query query;
  private final DocSet docSet;
  private final int docsToCollect;
  private final boolean needScores;
  private final String queryString;
  private final Query mainQuery;

  private TopDocsCollector topDocsCollector;
  private FilterCollector filterCollector;
  private MaxScoreCollector maxScoreCollector;
  private TopDocs topDocs;

  private QueryCommand(
      Sort sort,
      Query query,
      int docsToCollect,
      boolean needScores,
      DocSet docSet,
      String queryString,
      Query mainQuery) {
    this.sort = sort;
    this.query = query;
    this.docsToCollect = docsToCollect;
    this.needScores = needScores;
    this.docSet = docSet;
    this.queryString = queryString;
    this.mainQuery = mainQuery;
  }

  @Override
  public List 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(filterCollector);
  }

  @Override
  public void postCollect(IndexSearcher searcher) throws IOException {
    topDocs = topDocsCollector.topDocs();
    if (needScores) {
      // use mainQuery to populate the scores
      TopFieldCollector.populateScores(topDocs.scoreDocs, searcher, mainQuery);
    }
  }

  @Override
  public QueryCommandResult result() throws IOException {
    float maxScore;
    if (sort == null) {
      maxScore = topDocs.scoreDocs.length == 0 ? Float.NaN : topDocs.scoreDocs[0].score;
    } else {
      maxScore = maxScoreCollector == null ? Float.NaN : maxScoreCollector.getMaxScore();
    }
    return new QueryCommandResult(topDocs, filterCollector.getMatches(), maxScore);
  }

  @Override
  public String getKey() {
    return queryString != null ? queryString : query.toString();
  }

  @Override
  public Sort getGroupSort() {
    return sort;
  }

  @Override
  public Sort getWithinGroupSort() {
    return null;
  }
}