package org.apache.solr.search.grouping.distributed.requestfactory;
/*
* 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.ScoreDoc;
import org.apache.lucene.search.grouping.GroupDocs;
import org.apache.lucene.search.grouping.TopGroups;
import org.apache.solr.common.params.CommonParams;
import org.apache.solr.common.params.GroupParams;
import org.apache.solr.common.params.ModifiableSolrParams;
import org.apache.solr.common.params.ShardParams;
import org.apache.solr.common.util.StrUtils;
import org.apache.solr.handler.component.ResponseBuilder;
import org.apache.solr.handler.component.ShardDoc;
import org.apache.solr.handler.component.ShardRequest;
import org.apache.solr.schema.SchemaField;
import org.apache.solr.search.grouping.distributed.ShardRequestFactory;
import org.apache.solr.search.grouping.distributed.command.QueryCommandResult;
import java.util.*;
/**
*
*/
public class StoredFieldsShardRequestFactory implements ShardRequestFactory {
public ShardRequest[] constructRequest(ResponseBuilder rb) {
HashMap<String, Set<ShardDoc>> shardMap = new HashMap<String,Set<ShardDoc>>();
for (TopGroups<String> topGroups : rb.mergedTopGroups.values()) {
for (GroupDocs<String> group : topGroups.groups) {
mapShardToDocs(shardMap, group.scoreDocs);
}
}
for (QueryCommandResult queryCommandResult : rb.mergedQueryCommandResults.values()) {
mapShardToDocs(shardMap, queryCommandResult.getTopDocs().scoreDocs);
}
ShardRequest[] shardRequests = new ShardRequest[shardMap.size()];
SchemaField uniqueField = rb.req.getSchema().getUniqueKeyField();
int i = 0;
for (Collection<ShardDoc> shardDocs : shardMap.values()) {
ShardRequest sreq = new ShardRequest();
sreq.purpose = ShardRequest.PURPOSE_GET_FIELDS;
sreq.shards = new String[] {shardDocs.iterator().next().shard};
sreq.params = new ModifiableSolrParams();
sreq.params.add( rb.req.getParams());
sreq.params.remove(GroupParams.GROUP);
sreq.params.remove(CommonParams.SORT);
sreq.params.remove(ResponseBuilder.FIELD_SORT_VALUES);
String fl = sreq.params.get(CommonParams.FL);
if (fl != null) {
fl = fl.trim();
// currently, "score" is synonymous with "*,score" so
// don't add "id" if the fl is empty or "score" or it would change the meaning.
if (fl.length()!=0 && !"score".equals(fl) && !"*".equals(fl)) {
sreq.params.set(CommonParams.FL, fl+','+uniqueField.getName());
}
}
List<String> ids = new ArrayList<String>(shardDocs.size());
for (ShardDoc shardDoc : shardDocs) {
ids.add(shardDoc.id.toString());
}
sreq.params.add(ShardParams.IDS, StrUtils.join(ids, ','));
shardRequests[i++] = sreq;
}
return shardRequests;
}
private void mapShardToDocs(HashMap<String, Set<ShardDoc>> shardMap, ScoreDoc[] scoreDocs) {
for (ScoreDoc scoreDoc : scoreDocs) {
ShardDoc solrDoc = (ShardDoc) scoreDoc;
Set<ShardDoc> shardDocs = shardMap.get(solrDoc.shard);
if (shardDocs == null) {
shardMap.put(solrDoc.shard, shardDocs = new HashSet<ShardDoc>());
}
shardDocs.add(solrDoc);
}
}
}