org.apache.solr.search.grouping.distributed.command.TopGroupsFieldCommand 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.lang.reflect.Array;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import org.apache.lucene.queries.function.ValueSource;
import org.apache.lucene.search.Collector;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.Sort;
import org.apache.lucene.search.TopFieldCollector;
import org.apache.lucene.search.grouping.GroupDocs;
import org.apache.lucene.search.grouping.SearchGroup;
import org.apache.lucene.search.grouping.TermGroupSelector;
import org.apache.lucene.search.grouping.TopGroups;
import org.apache.lucene.search.grouping.TopGroupsCollector;
import org.apache.lucene.search.grouping.ValueSourceGroupSelector;
import org.apache.lucene.util.BytesRef;
import org.apache.lucene.util.mutable.MutableValue;
import org.apache.solr.schema.FieldType;
import org.apache.solr.schema.SchemaField;
import org.apache.solr.search.grouping.Command;

/**
 * Defines all collectors for retrieving the second phase and how to handle the collector result.
 */
public class TopGroupsFieldCommand implements Command> {

  public static class Builder {

    private Query query;
    private SchemaField field;
    private Sort groupSort;
    private Sort withinGroupSort;
    private Collection> firstPhaseGroups;
    private Integer maxDocPerGroup;
    private boolean needScores = false;
    private boolean needMaxScore = false;

    public Builder setQuery(Query query) {
      this.query = query;
      return this;
    }

    public Builder setField(SchemaField field) {
      this.field = field;
      return this;
    }

    public Builder setGroupSort(Sort groupSort) {
      this.groupSort = groupSort;
      return this;
    }

    public Builder setSortWithinGroup(Sort withinGroupSort) {
      this.withinGroupSort = withinGroupSort;
      return this;
    }

    public Builder setFirstPhaseGroups(Collection> firstPhaseGroups) {
      this.firstPhaseGroups = firstPhaseGroups;
      return this;
    }

    public Builder setMaxDocPerGroup(int maxDocPerGroup) {
      this.maxDocPerGroup = maxDocPerGroup;
      return this;
    }

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

    public Builder setNeedMaxScore(Boolean needMaxScore) {
      this.needMaxScore = needMaxScore;
      return this;
    }

    public TopGroupsFieldCommand build() {
      if (query == null
          || field == null
          || groupSort == null
          || withinGroupSort == null
          || firstPhaseGroups == null
          || maxDocPerGroup == null) {
        throw new IllegalStateException("All required fields must be set");
      }

      return new TopGroupsFieldCommand(
          query,
          field,
          groupSort,
          withinGroupSort,
          firstPhaseGroups,
          maxDocPerGroup,
          needScores,
          needMaxScore);
    }
  }

  private final Query query;
  private final SchemaField field;
  private final Sort groupSort;
  private final Sort withinGroupSort;
  private final Collection> firstPhaseGroups;
  private final int maxDocPerGroup;
  private final boolean needScores;
  private final boolean needMaxScore;
  private TopGroupsCollector secondPassCollector;
  private TopGroups topGroups;

  private TopGroupsFieldCommand(
      Query query,
      SchemaField field,
      Sort groupSort,
      Sort withinGroupSort,
      Collection> firstPhaseGroups,
      int maxDocPerGroup,
      boolean needScores,
      boolean needMaxScore) {
    this.query = query;
    this.field = field;
    this.groupSort = groupSort;
    this.withinGroupSort = withinGroupSort;
    this.firstPhaseGroups = firstPhaseGroups;
    this.maxDocPerGroup = maxDocPerGroup;
    this.needScores = needScores;
    this.needMaxScore = needMaxScore;
  }

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

    final List collectors = new ArrayList<>(1);
    final FieldType fieldType = field.getType();
    if (fieldType.getNumberType() != null) {
      ValueSource vs = fieldType.getValueSource(field, null);
      Collection> 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;
  }

  @Override
  @SuppressWarnings("unchecked")
  public void postCollect(IndexSearcher searcher) throws IOException {
    if (firstPhaseGroups.isEmpty()) {
      GroupDocs[] groups = (GroupDocs[]) Array.newInstance(GroupDocs.class, 0);
      topGroups =
          new TopGroups<>(groupSort.getSort(), withinGroupSort.getSort(), 0, 0, groups, Float.NaN);
      return;
    }

    FieldType fieldType = field.getType();
    TopGroups tg = secondPassCollector.getTopGroups(0);
    if (fieldType.getNumberType() != null) {
      topGroups = GroupConverter.fromMutable(field, (TopGroups) tg);
    } else {
      topGroups = (TopGroups) tg;
    }
    if (needScores) {
      for (GroupDocs group : topGroups.groups) {
        TopFieldCollector.populateScores(group.scoreDocs, searcher, query);
      }
    }
  }

  @Override
  public TopGroups result() throws IOException {
    return topGroups;
  }

  @Override
  public String getKey() {
    return field.getName();
  }

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

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