← Back to branch summary

~slub.team/goobi-indexserver/3.x

« back to all changes in this revision

Viewing changes to solr/core/src/java/org/apache/solr/search/grouping/distributed/shardresultserializer/TopGroupsResultTransformer.java

  • Committer: Sebastian Meyer
  • Date: 2012-08-03 09:12:40 UTC
  • Revision ID: sebastian.meyer@slub-dresden.de-20120803091240-x6861b0vabq1xror
https://launchpad.net/bugs/985487
Remove Lucene and Solr source code and add patches instead
Fix Bug #985487: Auto-suggestion for the search interface

Show diffs side-by-side

added added

removed removed

Lines of Context:
solr/core/src/java/org/apache/solr/search/grouping/distributed/shardresultserializer/TopGroupsResultTransformer.java
1
 
package org.apache.solr.search.grouping.distributed.shardresultserializer;
2
 
 
3
 
/*
4
 
 * Licensed to the Apache Software Foundation (ASF) under one or more
5
 
 * contributor license agreements.  See the NOTICE file distributed with
6
 
 * this work for additional information regarding copyright ownership.
7
 
 * The ASF licenses this file to You under the Apache License, Version 2.0
8
 
 * (the "License"); you may not use this file except in compliance with
9
 
 * the License.  You may obtain a copy of the License at
10
 
 *
11
 
 *     http://www.apache.org/licenses/LICENSE-2.0
12
 
 *
13
 
 * Unless required by applicable law or agreed to in writing, software
14
 
 * distributed under the License is distributed on an "AS IS" BASIS,
15
 
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
16
 
 * See the License for the specific language governing permissions and
17
 
 * limitations under the License.
18
 
 */
19
 
 
20
 
import org.apache.lucene.document.Document;
21
 
import org.apache.lucene.document.FieldSelector;
22
 
import org.apache.lucene.document.FieldSelectorResult;
23
 
import org.apache.lucene.search.FieldDoc;
24
 
import org.apache.lucene.search.ScoreDoc;
25
 
import org.apache.lucene.search.Sort;
26
 
import org.apache.lucene.search.TopDocs;
27
 
import org.apache.lucene.search.grouping.GroupDocs;
28
 
import org.apache.lucene.search.grouping.TopGroups;
29
 
import org.apache.solr.common.util.NamedList;
30
 
import org.apache.solr.handler.component.ResponseBuilder;
31
 
import org.apache.solr.handler.component.ShardDoc;
32
 
import org.apache.solr.schema.FieldType;
33
 
import org.apache.solr.schema.SchemaField;
34
 
import org.apache.solr.search.grouping.Command;
35
 
import org.apache.solr.search.grouping.distributed.command.QueryCommand;
36
 
import org.apache.solr.search.grouping.distributed.command.QueryCommandResult;
37
 
import org.apache.solr.search.grouping.distributed.command.TopGroupsFieldCommand;
38
 
 
39
 
import java.io.IOException;
40
 
import java.util.ArrayList;
41
 
import java.util.HashMap;
42
 
import java.util.List;
43
 
import java.util.Map;
44
 
 
45
 
/**
46
 
 * Implementation for transforming {@link TopGroups} and {@link TopDocs} into a {@link NamedList} structure and
47
 
 * visa versa.
48
 
 */
49
 
public class TopGroupsResultTransformer implements ShardResultTransformer<List<Command>, Map<String, ?>> {
50
 
 
51
 
  private final ResponseBuilder rb;
52
 
 
53
 
  public TopGroupsResultTransformer(ResponseBuilder rb) {
54
 
    this.rb = rb;
55
 
  }
56
 
 
57
 
  /**
58
 
   * {@inheritDoc}
59
 
   */
60
 
  public NamedList transform(List<Command> data) throws IOException {
61
 
    NamedList<NamedList> result = new NamedList<NamedList>();
62
 
    for (Command command : data) {
63
 
      NamedList commandResult;
64
 
      if (TopGroupsFieldCommand.class.isInstance(command)) {
65
 
        TopGroupsFieldCommand fieldCommand = (TopGroupsFieldCommand) command;
66
 
        SchemaField groupField = rb.req.getSearcher().getSchema().getField(fieldCommand.getKey());
67
 
        commandResult = serializeTopGroups(fieldCommand.result(), groupField);
68
 
      } else if (QueryCommand.class.isInstance(command)) {
69
 
        QueryCommand queryCommand = (QueryCommand) command;
70
 
        commandResult = serializeTopDocs(queryCommand.result());
71
 
      } else {
72
 
        commandResult = null;
73
 
      }
74
 
 
75
 
      result.add(command.getKey(), commandResult);
76
 
    }
77
 
    return result;
78
 
  }
79
 
 
80
 
  /**
81
 
   * {@inheritDoc}
82
 
   */
83
 
  public Map<String, ?> transformToNative(NamedList<NamedList> shardResponse, Sort groupSort, Sort sortWithinGroup, String shard) {
84
 
    Map<String, Object> result = new HashMap<String, Object>();
85
 
 
86
 
    for (Map.Entry<String, NamedList> entry : shardResponse) {
87
 
      String key = entry.getKey();
88
 
      NamedList commandResult = entry.getValue();
89
 
      Integer totalGroupedHitCount = (Integer) commandResult.get("totalGroupedHitCount");
90
 
      Integer totalHits = (Integer) commandResult.get("totalHits");
91
 
      if (totalHits != null) {
92
 
        Integer matches = (Integer) commandResult.get("matches");
93
 
        Float maxScore = (Float) commandResult.get("maxScore");
94
 
        if (maxScore == null) {
95
 
          maxScore = Float.NaN;
96
 
        }
97
 
 
98
 
        @SuppressWarnings("unchecked")
99
 
        List<NamedList<Object>> documents = (List<NamedList<Object>>) commandResult.get("documents");
100
 
        ScoreDoc[] scoreDocs = new ScoreDoc[documents.size()];
101
 
        int j = 0;
102
 
        for (NamedList<Object> document : documents) {
103
 
          Object uniqueId = document.get("id").toString();
104
 
          Float score = (Float) document.get("score");
105
 
          if (score == null) {
106
 
            score = Float.NaN;
107
 
          }
108
 
          Object[] sortValues = ((List) document.get("sortValues")).toArray();
109
 
          scoreDocs[j++] = new ShardDoc(score, sortValues, uniqueId, shard);
110
 
        }
111
 
        result.put(key, new QueryCommandResult(new TopDocs(totalHits, scoreDocs, maxScore), matches));
112
 
        continue;
113
 
      }
114
 
 
115
 
      Integer totalHitCount = (Integer) commandResult.get("totalHitCount");
116
 
      Integer totalGroupCount = (Integer) commandResult.get("totalGroupCount");
117
 
 
118
 
      List<GroupDocs<String>> groupDocs = new ArrayList<GroupDocs<String>>();
119
 
      for (int i = totalGroupCount == null ? 2 : 3; i < commandResult.size(); i++) {
120
 
        String groupValue = commandResult.getName(i);
121
 
        @SuppressWarnings("unchecked")
122
 
        NamedList<Object> groupResult = (NamedList<Object>) commandResult.getVal(i);
123
 
        Integer totalGroupHits = (Integer) groupResult.get("totalHits");
124
 
        Float maxScore = (Float) groupResult.get("maxScore");
125
 
        if (maxScore == null) {
126
 
          maxScore = Float.NaN;
127
 
        }
128
 
 
129
 
        @SuppressWarnings("unchecked")
130
 
        List<NamedList<Object>> documents = (List<NamedList<Object>>) groupResult.get("documents");
131
 
        ScoreDoc[] scoreDocs = new ScoreDoc[documents.size()];
132
 
        int j = 0;
133
 
        for (NamedList<Object> document : documents) {
134
 
          Object uniqueId = document.get("id").toString();
135
 
          Float score = (Float) document.get("score");
136
 
          if (score == null) {
137
 
            score = Float.NaN;
138
 
          }
139
 
          Object[] sortValues = ((List) document.get("sortValues")).toArray();
140
 
          scoreDocs[j++] = new ShardDoc(score, sortValues, uniqueId, shard);
141
 
        }
142
 
 
143
 
        String groupValueRef = groupValue != null ? groupValue : null;
144
 
        groupDocs.add(new GroupDocs<String>(maxScore, totalGroupHits, scoreDocs, groupValueRef, null));
145
 
      }
146
 
 
147
 
      @SuppressWarnings("unchecked")
148
 
      GroupDocs<String>[] groupDocsArr = groupDocs.toArray(new GroupDocs[groupDocs.size()]);
149
 
      TopGroups<String> topGroups = new TopGroups<String>(
150
 
        groupSort.getSort(), sortWithinGroup.getSort(), totalHitCount, totalGroupedHitCount, groupDocsArr
151
 
      );
152
 
      if (totalGroupCount != null) {
153
 
        topGroups = new TopGroups<String>(topGroups, totalGroupCount);
154
 
      }
155
 
 
156
 
      result.put(key, topGroups);
157
 
    }
158
 
 
159
 
    return result;
160
 
  }
161
 
 
162
 
  protected NamedList serializeTopGroups(TopGroups<String> data, SchemaField groupField) throws IOException {
163
 
    NamedList<Object> result = new NamedList<Object>();
164
 
    result.add("totalGroupedHitCount", data.totalGroupedHitCount);
165
 
    result.add("totalHitCount", data.totalHitCount);
166
 
    if (data.totalGroupCount != null) {
167
 
      result.add("totalGroupCount", data.totalGroupCount);
168
 
    }
169
 
    SchemaField uniqueField = rb.req.getSearcher().getSchema().getUniqueKeyField();
170
 
    for (GroupDocs<String> searchGroup : data.groups) {
171
 
      NamedList<Object> groupResult = new NamedList<Object>();
172
 
      groupResult.add("totalHits", searchGroup.totalHits);
173
 
      if (!Float.isNaN(searchGroup.maxScore)) {
174
 
        groupResult.add("maxScore", searchGroup.maxScore);
175
 
      }
176
 
 
177
 
      List<NamedList<Object>> documents = new ArrayList<NamedList<Object>>();
178
 
      for (int i = 0; i < searchGroup.scoreDocs.length; i++) {
179
 
        NamedList<Object> document = new NamedList<Object>();
180
 
        documents.add(document);
181
 
 
182
 
        Document doc = retrieveDocument(uniqueField, searchGroup.scoreDocs[i].doc);
183
 
        document.add("id", uniqueField.getType().toObject(doc.getFieldable(uniqueField.getName())));
184
 
        if (!Float.isNaN(searchGroup.scoreDocs[i].score))  {
185
 
          document.add("score", searchGroup.scoreDocs[i].score);
186
 
        }
187
 
        if (!(searchGroup.scoreDocs[i] instanceof FieldDoc)) {
188
 
          continue;
189
 
        }
190
 
 
191
 
        FieldDoc fieldDoc = (FieldDoc) searchGroup.scoreDocs[i];
192
 
        Object[] convertedSortValues  = new Object[fieldDoc.fields.length];
193
 
        for (int j = 0; j < fieldDoc.fields.length; j++) {
194
 
          Object sortValue  = fieldDoc.fields[j];
195
 
          Sort sortWithinGroup = rb.getGroupingSpec().getSortWithinGroup();
196
 
          SchemaField field = sortWithinGroup.getSort()[j].getField() != null ? rb.req.getSearcher().getSchema().getFieldOrNull(sortWithinGroup.getSort()[j].getField()) : null;
197
 
          if (field != null) {
198
 
            FieldType fieldType = field.getType();
199
 
            if (sortValue instanceof String) {
200
 
              sortValue = fieldType.toObject(field.createField(fieldType.indexedToReadable((String) sortValue), 0.0f));
201
 
            }
202
 
          }
203
 
          convertedSortValues[j] = sortValue;
204
 
        }
205
 
        document.add("sortValues", convertedSortValues);
206
 
      }
207
 
      groupResult.add("documents", documents);
208
 
      String groupValue = searchGroup.groupValue != null ? groupField.getType().indexedToReadable(searchGroup.groupValue): null;
209
 
      result.add(groupValue, groupResult);
210
 
    }
211
 
 
212
 
    return result;
213
 
  }
214
 
 
215
 
  protected NamedList serializeTopDocs(QueryCommandResult result) throws IOException {
216
 
    NamedList<Object> queryResult = new NamedList<Object>();
217
 
    queryResult.add("matches", result.getMatches());
218
 
    queryResult.add("totalHits", result.getTopDocs().totalHits);
219
 
    if (rb.getGroupingSpec().isNeedScore()) {
220
 
      queryResult.add("maxScore", result.getTopDocs().getMaxScore());
221
 
    }
222
 
    List<NamedList> documents = new ArrayList<NamedList>();
223
 
    queryResult.add("documents", documents);
224
 
 
225
 
    SchemaField uniqueField = rb.req.getSearcher().getSchema().getUniqueKeyField();
226
 
    for (ScoreDoc scoreDoc : result.getTopDocs().scoreDocs) {
227
 
      NamedList<Object> document = new NamedList<Object>();
228
 
      documents.add(document);
229
 
 
230
 
      Document doc = retrieveDocument(uniqueField, scoreDoc.doc);
231
 
      document.add("id", uniqueField.getType().toObject(doc.getFieldable(uniqueField.getName())));
232
 
      if (rb.getGroupingSpec().isNeedScore())  {
233
 
        document.add("score", scoreDoc.score);
234
 
      }
235
 
      if (!FieldDoc.class.isInstance(scoreDoc)) {
236
 
        continue;
237
 
      }
238
 
 
239
 
      FieldDoc fieldDoc = (FieldDoc) scoreDoc;
240
 
      Object[] convertedSortValues  = new Object[fieldDoc.fields.length];
241
 
      for (int j = 0; j < fieldDoc.fields.length; j++) {
242
 
        Object sortValue  = fieldDoc.fields[j];
243
 
        Sort groupSort = rb.getGroupingSpec().getGroupSort();
244
 
        SchemaField field = groupSort.getSort()[j].getField() != null ? rb.req.getSearcher().getSchema().getFieldOrNull(groupSort.getSort()[j].getField()) : null;
245
 
        if (field != null) {
246
 
          FieldType fieldType = field.getType();
247
 
          if (sortValue instanceof String) {
248
 
            sortValue = fieldType.toObject(field.createField(fieldType.indexedToReadable((String) sortValue), 0.0f));
249
 
          }
250
 
        }
251
 
        convertedSortValues[j] = sortValue;
252
 
      }
253
 
      document.add("sortValues", convertedSortValues);
254
 
    }
255
 
 
256
 
    return queryResult;
257
 
  }
258
 
 
259
 
  private Document retrieveDocument(final SchemaField uniqueField, int doc) throws IOException {
260
 
    FieldSelector fieldSelectorVisitor = new FieldSelector() {
261
 
 
262
 
      public FieldSelectorResult accept(String fieldName) {
263
 
        if (uniqueField.getName().equals(fieldName)) {
264
 
          return FieldSelectorResult.LOAD_AND_BREAK;
265
 
        }
266
 
        return FieldSelectorResult.NO_LOAD;
267
 
      }
268
 
    };
269
 
    return rb.req.getSearcher().doc(doc, fieldSelectorVisitor);
270
 
  }
271
 
 
272
 
}