| package org.apache.solr.search.grouping.distributed.command; |
| |
| /* |
| * 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. |
| */ |
| |
| import org.apache.lucene.search.Collector; |
| import org.apache.lucene.search.Sort; |
| import org.apache.lucene.search.grouping.*; |
| import org.apache.solr.schema.SchemaField; |
| import org.apache.solr.search.grouping.Command; |
| |
| import java.io.IOException; |
| import java.util.ArrayList; |
| import java.util.Collection; |
| import java.util.Collections; |
| import java.util.List; |
| |
| /** |
| * |
| */ |
| public class TopGroupsFieldCommand implements Command<TopGroups<String>> { |
| |
| public static class Builder { |
| |
| private SchemaField field; |
| private Sort groupSort; |
| private Sort sortWithinGroup; |
| private Collection<SearchGroup<String>> firstPhaseGroups; |
| private Integer maxDocPerGroup; |
| private boolean needScores = false; |
| private boolean needMaxScore = false; |
| private boolean needGroupCount = false; |
| |
| public Builder setField(SchemaField field) { |
| this.field = field; |
| return this; |
| } |
| |
| public Builder setGroupSort(Sort groupSort) { |
| this.groupSort = groupSort; |
| return this; |
| } |
| |
| public Builder setSortWithinGroup(Sort sortWithinGroup) { |
| this.sortWithinGroup = sortWithinGroup; |
| return this; |
| } |
| |
| public Builder setFirstPhaseGroups(Collection<SearchGroup<String>> 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 Builder setNeedGroupCount(Boolean needGroupCount) { |
| this.needGroupCount = needGroupCount; |
| return this; |
| } |
| |
| public TopGroupsFieldCommand build() { |
| if (field == null || groupSort == null || sortWithinGroup == null || firstPhaseGroups == null || |
| maxDocPerGroup == null) { |
| throw new IllegalStateException("All required fields must be set"); |
| } |
| |
| return new TopGroupsFieldCommand(field, groupSort, sortWithinGroup, firstPhaseGroups, maxDocPerGroup, needScores, needMaxScore, needGroupCount); |
| } |
| |
| } |
| |
| private final SchemaField field; |
| private final Sort groupSort; |
| private final Sort sortWithinGroup; |
| private final Collection<SearchGroup<String>> firstPhaseGroups; |
| private final int maxDocPerGroup; |
| private final boolean needScores; |
| private final boolean needMaxScore; |
| private final boolean needGroupCount; |
| |
| private TermSecondPassGroupingCollector secondPassCollector; |
| private TermAllGroupsCollector allGroupsCollector; |
| |
| private TopGroupsFieldCommand(SchemaField field, |
| Sort groupSort, |
| Sort sortWithinGroup, |
| Collection<SearchGroup<String>> firstPhaseGroups, |
| int maxDocPerGroup, |
| boolean needScores, |
| boolean needMaxScore, |
| boolean needGroupCount) { |
| this.field = field; |
| this.groupSort = groupSort; |
| this.sortWithinGroup = sortWithinGroup; |
| this.firstPhaseGroups = firstPhaseGroups; |
| this.maxDocPerGroup = maxDocPerGroup; |
| this.needScores = needScores; |
| this.needMaxScore = needMaxScore; |
| this.needGroupCount = needGroupCount; |
| } |
| |
| public List<Collector> create() throws IOException { |
| if (firstPhaseGroups.isEmpty()) { |
| return Collections.emptyList(); |
| } |
| |
| List<Collector> collectors = new ArrayList<Collector>(); |
| secondPassCollector = new TermSecondPassGroupingCollector( |
| field.getName(), firstPhaseGroups, groupSort, sortWithinGroup, maxDocPerGroup, needScores, needMaxScore, true |
| ); |
| collectors.add(secondPassCollector); |
| if (!needGroupCount) { |
| return collectors; |
| } |
| allGroupsCollector = new TermAllGroupsCollector(field.getName()); |
| collectors.add(allGroupsCollector); |
| return collectors; |
| } |
| |
| @SuppressWarnings("unchecked") |
| public TopGroups<String> result() { |
| if (firstPhaseGroups.isEmpty()) { |
| return new TopGroups<String>(groupSort.getSort(), sortWithinGroup.getSort(), 0, 0, new GroupDocs[0]); |
| } |
| |
| TopGroups<String> result = secondPassCollector.getTopGroups(0); |
| if (allGroupsCollector != null) { |
| result = new TopGroups<String>(result, allGroupsCollector.getGroupCount()); |
| } |
| return result; |
| } |
| |
| public String getKey() { |
| return field.getName(); |
| } |
| |
| public Sort getGroupSort() { |
| return groupSort; |
| } |
| |
| public Sort getSortWithinGroup() { |
| return sortWithinGroup; |
| } |
| } |