org.apache.lucene.search.grouping.TopGroups Java Examples

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

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

@SuppressWarnings({"unchecked", "rawtypes"})
static TopGroups<BytesRef> fromMutable(SchemaField field, TopGroups<MutableValue> values) {
  if (values == null) {
    return null;
  }
  
  FieldType fieldType = field.getType();

  GroupDocs<BytesRef> groupDocs[] = new GroupDocs[values.groups.length];

  for (int i = 0; i < values.groups.length; i++) {
    GroupDocs<MutableValue> original = values.groups[i];
    final BytesRef groupValue;
    if (original.groupValue.exists) {
      BytesRefBuilder binary = new BytesRefBuilder();
      fieldType.readableToIndexed(Utils.OBJECT_TO_STRING.apply(original.groupValue.toObject()), binary);
      groupValue = binary.get();
    } else {
      groupValue = null;
    }
    groupDocs[i] = new GroupDocs<>(original.score, original.maxScore, original.totalHits, original.scoreDocs, groupValue, original.groupSortValues);
  }
  
  return new TopGroups<>(values.groupSort, values.withinGroupSort, values.totalHitCount, values.totalGroupedHitCount, groupDocs, values.maxScore);
}
 

public Map<DocumentType, List<SearchResult>> search(String searchString) throws ParseException {
    Map<DocumentType, List<SearchResult>> resultMap = new TreeMap<DocumentType, List<SearchResult>>();
    try {
        Query query = parser.parse(searchString);
        final SecondPassGroupingCollector collector = new SecondPassGroupingCollector("documentType", searchGroups,
                Sort.RELEVANCE, null, 5, true, false, true);
        searcher.search(query, collector);
        final TopGroups groups = collector.getTopGroups(0);
        for (GroupDocs groupDocs : groups.groups) {
            DocumentType docType = DocumentType.valueOf(groupDocs.groupValue);
            List<SearchResult> results = new ArrayList<SearchResult>();
            for (ScoreDoc scoreDoc : groupDocs.scoreDocs) {
                Document doc = searcher.doc(scoreDoc.doc);
                SearchResult result = new SearchResult(
                        docType,
                        doc.get("name"),
                        doc.get("url"),
                        doc.get("className"),
                        doc.get("package"),
                        doc.get("ensemblePath"),
                        doc.get("shortDescription")
                );
                results.add(result);
            }
            resultMap.put(docType, results);
        }
    } catch (IOException e) {
        e.printStackTrace();
    }
    return resultMap;
}
 

public Map<DocumentType, List<SearchResult>> search(String searchString) throws ParseException {
    Map<DocumentType, List<SearchResult>> resultMap = new TreeMap<DocumentType, List<SearchResult>>();
    try {
        Query query = parser.parse(searchString);
        final SecondPassGroupingCollector collector = new SecondPassGroupingCollector("documentType", searchGroups,
                Sort.RELEVANCE, null, 5, true, false, true);
        searcher.search(query, collector);
        final TopGroups groups = collector.getTopGroups(0);
        for (GroupDocs groupDocs : groups.groups) {
            DocumentType docType = DocumentType.valueOf(groupDocs.groupValue);
            List<SearchResult> results = new ArrayList<SearchResult>();
            for (ScoreDoc scoreDoc : groupDocs.scoreDocs) {
                Document doc = searcher.doc(scoreDoc.doc);
                SearchResult result = new SearchResult(
                        docType,
                        doc.get("name"),
                        doc.get("url"),
                        doc.get("className"),
                        doc.get("package"),
                        doc.get("ensemblePath"),
                        doc.get("shortDescription")
                );
                results.add(result);
            }
            resultMap.put(docType, results);
        }
    } catch (IOException e) {
        e.printStackTrace();
    }
    return resultMap;
}
 

public LindenResult parse(TopDocs topDocs, TopGroups<TopDocs> topGroupedDocs,
                          Facets facets, FacetsCollector facetsCollector) throws IOException {
  LindenResult result = new LindenResult();
  List<LindenHit> lindenHits;
  int totalHits = 0;
  if (topDocs != null) {
    totalHits = topDocs.totalHits;
    lindenHits = parseLindenHits(topDocs.scoreDocs);
  } else if (topGroupedDocs != null) {
    lindenHits = new ArrayList<>();
    totalHits = topGroupedDocs.totalHitCount;
    for (GroupDocs<TopDocs> group : topGroupedDocs.groups) {
      List<LindenHit> groupHits = parseLindenHits(group.scoreDocs);
      LindenHit hitGroup = new LindenHit(groupHits.get(0)).setGroupHits(groupHits);
      String groupField = request.getGroupParam().getGroupField();
      String groupValue = LindenUtil.getFieldStringValue(leaves, group.scoreDocs[0].doc, groupField);
      if (!hitGroup.isSetFields()) {
        hitGroup.setFields(new HashMap<String, String>());
      }
      hitGroup.getFields().put(groupField, groupValue);
      lindenHits.add(hitGroup);
    }
    int groupTotal = topGroupedDocs.totalGroupCount == null ? 0 : topGroupedDocs.totalGroupCount;
    result.setTotalGroups(groupTotal);
    result.setTotalGroupHits(topGroupedDocs.totalGroupedHitCount);
  } else {
    lindenHits = new ArrayList<>();
  }
  result.setTotalHits(totalHits);
  result.setHits(lindenHits);
  parseFacets(result, facets, facetsCollector);
  result.setQueryInfo(new QueryInfo().setQuery(query.toString()));
  if (filter != null) {
    result.getQueryInfo().setFilter(filter.toString());
  }
  if (sort != null) {
    result.getQueryInfo().setSort(sort.toString());
  }
  return result;
}
 

@SuppressWarnings({"unchecked", "rawtypes"})
protected void groupedFinishStage(final ResponseBuilder rb) {
  // To have same response as non-distributed request.
  GroupingSpecification groupSpec = rb.getGroupingSpec();
  if (rb.mergedTopGroups.isEmpty()) {
    for (String field : groupSpec.getFields()) {
      rb.mergedTopGroups.put(field, new TopGroups(null, null, 0, 0, new GroupDocs[]{}, Float.NaN));
    }
    rb.resultIds = new HashMap<>();
  }

  EndResultTransformer.SolrDocumentSource solrDocumentSource = doc -> {
    ShardDoc solrDoc = (ShardDoc) doc;
    return rb.retrievedDocuments.get(solrDoc.id);
  };
  EndResultTransformer endResultTransformer;
  if (groupSpec.isMain()) {
    endResultTransformer = MAIN_END_RESULT_TRANSFORMER;
  } else if (Grouping.Format.grouped == groupSpec.getResponseFormat()) {
    endResultTransformer = new GroupedEndResultTransformer(rb.req.getSearcher());
  } else if (Grouping.Format.simple == groupSpec.getResponseFormat() && !groupSpec.isMain()) {
    endResultTransformer = SIMPLE_END_RESULT_TRANSFORMER;
  } else {
    return;
  }
  Map<String, Object> combinedMap = new LinkedHashMap<>();
  combinedMap.putAll(rb.mergedTopGroups);
  combinedMap.putAll(rb.mergedQueryCommandResults);
  endResultTransformer.transform(combinedMap, rb, solrDocumentSource);
}
 

protected NamedList serializeTopGroups(TopGroups<BytesRef> data, SchemaField groupField) throws IOException {
  NamedList<Object> result = new NamedList<>();
  result.add("totalGroupedHitCount", data.totalGroupedHitCount);
  result.add("totalHitCount", data.totalHitCount);
  if (data.totalGroupCount != null) {
    result.add("totalGroupCount", data.totalGroupCount);
  }

  final IndexSchema schema = rb.req.getSearcher().getSchema();
  SchemaField uniqueField = schema.getUniqueKeyField();
  for (GroupDocs<BytesRef> searchGroup : data.groups) {
    NamedList<Object> groupResult = new NamedList<>();
    assert searchGroup.totalHits.relation == TotalHits.Relation.EQUAL_TO;
    groupResult.add("totalHits", searchGroup.totalHits.value);
    if (!Float.isNaN(searchGroup.maxScore)) {
      groupResult.add("maxScore", searchGroup.maxScore);
    }

    List<NamedList<Object>> documents = new ArrayList<>();
    for (int i = 0; i < searchGroup.scoreDocs.length; i++) {
      NamedList<Object> document = new NamedList<>();
      documents.add(document);

      Document doc = retrieveDocument(uniqueField, searchGroup.scoreDocs[i].doc);
      document.add(ID, uniqueField.getType().toExternal(doc.getField(uniqueField.getName())));
      if (!Float.isNaN(searchGroup.scoreDocs[i].score))  {
        document.add("score", searchGroup.scoreDocs[i].score);
      }
      if (!(searchGroup.scoreDocs[i] instanceof FieldDoc)) {
        continue; // thus don't add sortValues below
      }

      FieldDoc fieldDoc = (FieldDoc) searchGroup.scoreDocs[i];
      Object[] convertedSortValues  = new Object[fieldDoc.fields.length];
      for (int j = 0; j < fieldDoc.fields.length; j++) {
        Object sortValue  = fieldDoc.fields[j];
        Sort withinGroupSort = rb.getGroupingSpec().getWithinGroupSortSpec().getSort();
        SchemaField field = withinGroupSort.getSort()[j].getField() != null ? schema.getFieldOrNull(withinGroupSort.getSort()[j].getField()) : null;
        if (field != null) {
          FieldType fieldType = field.getType();
          if (sortValue != null) {
            sortValue = fieldType.marshalSortValue(sortValue);
          }
        }
        convertedSortValues[j] = sortValue;
      }
      document.add("sortValues", convertedSortValues);
    }
    groupResult.add("documents", documents);
    String groupValue = searchGroup.groupValue != null ?
        groupField.getType().indexedToReadable(searchGroup.groupValue, new CharsRefBuilder()).toString(): null;
    result.add(groupValue, groupResult);
  }

  return result;
}
 

@Override
public ShardRequest[] constructRequest(ResponseBuilder rb) {
  HashMap<String, Set<ShardDoc>> shardMap = new HashMap<>();
  for (TopGroups<BytesRef> topGroups : rb.mergedTopGroups.values()) {
    for (GroupDocs<BytesRef> 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);
    
    // we need to ensure the uniqueField is included for collating docs with their return fields
    if (! rb.rsp.getReturnFields().wantsField(uniqueField.getName())) {
      // the user didn't ask for it, so we have to...
      sreq.params.add(CommonParams.FL, uniqueField.getName());
    }

    List<String> ids = new ArrayList<>(shardDocs.size());
    for (ShardDoc shardDoc : shardDocs) {
      ids.add(shardDoc.id.toString());
    }
    sreq.params.add(ShardParams.IDS, StrUtils.join(ids, ','));
    shardRequests[i++] = sreq;
  }

  return shardRequests;
}
 

@Override
@SuppressWarnings("unchecked")
public TopGroups<BytesRef> result() throws IOException {
  return topGroups;
}