Java Code Examples for org.apache.lucene.search.Sort#getSort()

@SuppressWarnings({"unchecked","rawtypes"})
private Comparable<?>[] fillFields(GroupDoc d, Sort sort) {
  final SortField[] sortFields = sort.getSort();
  final Comparable<?>[] fields = new Comparable[sortFields.length];
  for(int fieldIDX=0;fieldIDX<sortFields.length;fieldIDX++) {
    final Comparable<?> c;
    final SortField sf = sortFields[fieldIDX];
    if (sf.getType() == SortField.Type.SCORE) {
      c = d.score;
    } else if (sf.getField().equals("sort1")) {
      c = d.sort1;
    } else if (sf.getField().equals("sort2")) {
      c = d.sort2;
    } else {
      assertEquals("id", sf.getField());
      c = d.id;
    }
    fields[fieldIDX] = c;
  }
  return fields;
}
 

private SortField parse(String a) {
    final SolrQueryRequest req = req("q", "{!parent which=type_s1:parent}whatever_s1:foo",
        "q2", "{!parent which=type_s1:parent}nomater_s1:what",
        "notbjq", "foo_s1:bar");
    try {
    final SortSpec spec = SortSpecParsing.parseSortSpec(a,
        req);
    assertNull(spec.getSchemaFields().get(0));
    final Sort sort = spec.getSort();
    final SortField field = sort.getSort()[0];
    assertNotNull(field);
    return field;
  } finally {
    req.close();
  }
}
 

@Override
@SuppressWarnings({"rawtypes"})
public Map<String, SearchGroupsFieldCommandResult> transformToNative(NamedList<NamedList> shardResponse, Sort groupSort, Sort withinGroupSort, String shard) {
  final Map<String, SearchGroupsFieldCommandResult> result = new HashMap<>(shardResponse.size());
  for (Map.Entry<String, NamedList> command : shardResponse) {
    List<SearchGroup<BytesRef>> searchGroups = new ArrayList<>();
    NamedList topGroupsAndGroupCount = command.getValue();
    @SuppressWarnings({"unchecked"})
    final NamedList<List<Comparable>> rawSearchGroups = (NamedList<List<Comparable>>) topGroupsAndGroupCount.get(TOP_GROUPS);
    if (rawSearchGroups != null) {
      final SchemaField groupField = searcher.getSchema().getFieldOrNull(command.getKey());
      final SortField[] groupSortField = groupSort.getSort();
      for (Map.Entry<String, List<Comparable>> rawSearchGroup : rawSearchGroups){
        SearchGroup<BytesRef> searchGroup = deserializeOneSearchGroup(
            groupField, rawSearchGroup.getKey(),
            groupSortField, rawSearchGroup.getValue());
        searchGroups.add(searchGroup);
      }
    }

    final Integer groupCount = (Integer) topGroupsAndGroupCount.get(GROUP_COUNT);
    result.put(command.getKey(), new SearchGroupsFieldCommandResult(groupCount, searchGroups));
  }
  return result;
}
 

protected SortingGroupHead(Sort sort, T groupValue, int doc, LeafReaderContext context, Scorable scorer) throws IOException {
  super(groupValue, doc, context.docBase);
  final SortField[] sortFields = sort.getSort();
  comparators = new FieldComparator[sortFields.length];
  leafComparators = new LeafFieldComparator[sortFields.length];
  for (int i = 0; i < sortFields.length; i++) {
    comparators[i] = sortFields[i].getComparator(1, i);
    leafComparators[i] = comparators[i].getLeafComparator(context);
    leafComparators[i].setScorer(scorer);
    leafComparators[i].copy(0, doc);
    leafComparators[i].setBottom(0);
  }
}
 

/**
 * Create the single pass collector.
 *
 *  @param groupSort The {@link Sort} used to sort the
 *    groups.  The top sorted document within each group
 *    according to groupSort, determines how that group
 *    sorts against other groups.  This must be non-null,
 *    ie, if you want to groupSort by relevance use
 *    Sort.RELEVANCE.
 *  @param topNGroups How many top groups to keep.
 *  @param needsScores true if the collected documents
 *    require scores, either because relevance is included
 *    in the withinGroupSort or because you plan to pass true
 *    for either getSscores or getMaxScores to {@link
 *    #getTopGroups}
 *  @param lastDocPerGroup a {@link Weight} that marks the
 *    last document in each group.
 */
public BlockGroupingCollector(Sort groupSort, int topNGroups, boolean needsScores, Weight lastDocPerGroup) {

  if (topNGroups < 1) {
    throw new IllegalArgumentException("topNGroups must be >= 1 (got " + topNGroups + ")");
  }

  groupQueue = new GroupQueue(topNGroups);
  pendingSubDocs = new int[10];
  if (needsScores) {
    pendingSubScores = new float[10];
  }

  this.needsScores = needsScores;
  this.lastDocPerGroup = lastDocPerGroup;

  this.groupSort = groupSort;
  
  this.topNGroups = topNGroups;

  final SortField[] sortFields = groupSort.getSort();
  comparators = new FieldComparator<?>[sortFields.length];
  leafComparators = new LeafFieldComparator[sortFields.length];
  compIDXEnd = comparators.length - 1;
  reversed = new int[sortFields.length];
  for (int i = 0; i < sortFields.length; i++) {
    final SortField sortField = sortFields[i];
    comparators[i] = sortField.getComparator(topNGroups, i);
    reversed[i] = sortField.getReverse() ? -1 : 1;
  }
}
 

/**
 * Create the first pass collector.
 *
 * @param groupSelector a GroupSelector used to defined groups
 * @param groupSort The {@link Sort} used to sort the
 *    groups.  The top sorted document within each group
 *    according to groupSort, determines how that group
 *    sorts against other groups.  This must be non-null,
 *    ie, if you want to groupSort by relevance use
 *    Sort.RELEVANCE.
 * @param topNGroups How many top groups to keep.
 */
@SuppressWarnings({"unchecked", "rawtypes"})
public FirstPassGroupingCollector(GroupSelector<T> groupSelector, Sort groupSort, int topNGroups) {
  this.groupSelector = groupSelector;
  if (topNGroups < 1) {
    throw new IllegalArgumentException("topNGroups must be >= 1 (got " + topNGroups + ")");
  }

  // TODO: allow null groupSort to mean "by relevance",
  // and specialize it?

  this.topNGroups = topNGroups;
  this.needsScores = groupSort.needsScores();
  final SortField[] sortFields = groupSort.getSort();
  comparators = new FieldComparator[sortFields.length];
  leafComparators = new LeafFieldComparator[sortFields.length];
  compIDXEnd = comparators.length - 1;
  reversed = new int[sortFields.length];
  for (int i = 0; i < sortFields.length; i++) {
    final SortField sortField = sortFields[i];

    // use topNGroups + 1 so we have a spare slot to use for comparing (tracked by this.spareSlot):
    comparators[i] = sortField.getComparator(topNGroups + 1, i);
    reversed[i] = sortField.getReverse() ? -1 : 1;
  }

  spareSlot = topNGroups;
  groupMap = new HashMap<>(topNGroups);
}
 

@SuppressWarnings({"unchecked", "rawtypes"})
public GroupComparator(Sort groupSort) {
  final SortField[] sortFields = groupSort.getSort();
  comparators = new FieldComparator[sortFields.length];
  reversed = new int[sortFields.length];
  for (int compIDX = 0; compIDX < sortFields.length; compIDX++) {
    final SortField sortField = sortFields[compIDX];
    comparators[compIDX] = sortField.getComparator(1, compIDX);
    reversed[compIDX] = sortField.getReverse() ? -1 : 1;
  }
}
 

private Comparator<GroupDoc> getComparator(Sort sort) {
  final SortField[] sortFields = sort.getSort();
  return new Comparator<GroupDoc>() {
    @Override
    public int compare(GroupDoc d1, GroupDoc d2) {
      for(SortField sf : sortFields) {
        final int cmp;
        if (sf.getType() == SortField.Type.SCORE) {
          if (d1.score > d2.score) {
            cmp = -1;
          } else if (d1.score < d2.score) {
            cmp = 1;
          } else {
            cmp = 0;
          }
        } else if (sf.getField().equals("sort1")) {
          cmp = d1.sort1.compareTo(d2.sort1);
        } else if (sf.getField().equals("sort2")) {
          cmp = d1.sort2.compareTo(d2.sort2);
        } else {
          assertEquals(sf.getField(), "id");
          cmp = d1.id - d2.id;
        }
        if (cmp != 0) {
          return sf.getReverse() ? -cmp : cmp;
        }
      }
      // Our sort always fully tie breaks:
      fail();
      return 0;
    }
  };
}
 

private int getSortScoreFieldPos(Sort sort) {
  int sortScoreField = -1;
  if (sort != null) {
    for (int i = 0; i < sort.getSort().length; ++i) {
      if (sort.getSort()[i].getType() == SortField.Type.SCORE) {
        sortScoreField = i;
        break;
      }
    }
  }
  return sortScoreField;
}
 

private Comparator<GroupDoc> getComparator(Sort sort, final boolean sortByScoreOnly, final int[] fieldIdToDocID) {
  final SortField[] sortFields = sort.getSort();
  return new Comparator<GroupDoc>() {
    @Override
    public int compare(GroupDoc d1, GroupDoc d2) {
      for (SortField sf : sortFields) {
        final int cmp;
        if (sf.getType() == SortField.Type.SCORE) {
          if (d1.score > d2.score) {
            cmp = -1;
          } else if (d1.score < d2.score) {
            cmp = 1;
          } else {
            cmp = sortByScoreOnly ? fieldIdToDocID[d1.id] - fieldIdToDocID[d2.id] : 0;
          }
        } else if (sf.getField().equals("sort1")) {
          cmp = d1.sort1.compareTo(d2.sort1);
        } else if (sf.getField().equals("sort2")) {
          cmp = d1.sort2.compareTo(d2.sort2);
        } else if (sf.getField().equals("sort3")) {
          cmp = d1.sort3.compareTo(d2.sort3);
        } else {
          assertEquals(sf.getField(), "id");
          cmp = d1.id - d2.id;
        }
        if (cmp != 0) {
          return sf.getReverse() ? -cmp : cmp;
        }
      }
      // Our sort always fully tie breaks:
      fail();
      return 0;
    }
  };
}
 

/**
 * Set the {@link Sort} order to use for all (flushed and merged) segments.
 */
public IndexWriterConfig setIndexSort(Sort sort) {
  for (SortField sortField : sort.getSort()) {
    if (sortField.getIndexSorter() == null) {
      throw new IllegalArgumentException("Cannot sort index with sort field " + sortField);
    }
  }
  this.indexSort = sort;
  this.indexSortFields = Arrays.stream(sort.getSort()).map(SortField::getField).collect(Collectors.toSet());
  return this;
}
 

public static boolean includesScore(Sort sort) {
  if (sort==null) return true;
  for (SortField sf : sort.getSort()) {
    if (sf.getType() == SortField.Type.SCORE) return true;
  }
  return false;
}
 

private AllGroupHeadsCollector(GroupSelector<T> selector, Sort sort) {
  this.groupSelector = selector;
  this.sort = sort;
  this.reversed = new int[sort.getSort().length];
  final SortField[] sortFields = sort.getSort();
  for (int i = 0; i < sortFields.length; i++) {
    reversed[i] = sortFields[i].getReverse() ? -1 : 1;
  }
  this.compIDXEnd = this.reversed.length - 1;
}
 

/** Returns whether collection can be early-terminated if it sorts with the
 *  provided {@link Sort} and if segments are merged with the provided
 *  {@link Sort}. */
public static boolean canEarlyTerminate(Sort searchSort, Sort mergePolicySort) {
  final SortField[] fields1 = searchSort.getSort();
  final SortField[] fields2 = mergePolicySort.getSort();
  // early termination is possible if fields1 is a prefix of fields2
  if (fields1.length > fields2.length) {
    return false;
  }
  return Arrays.asList(fields1).equals(Arrays.asList(fields2).subList(0, fields1.length));
}
 

void assertQuery(Query query, Sort sort) throws Exception {
  int size = TestUtil.nextInt(random(), 1, searcher.getIndexReader().maxDoc() / 5);
  TopDocs expected = searcher.search(query, size, sort, random().nextBoolean());
  
  // make our actual sort, mutating original by replacing some of the 
  // sortfields with equivalent expressions
  
  SortField original[] = sort.getSort();
  SortField mutated[] = new SortField[original.length];
  for (int i = 0; i < mutated.length; i++) {
    if (random().nextInt(3) > 0) {
      SortField s = original[i];
      Expression expr = JavascriptCompiler.compile(s.getField());
      SimpleBindings simpleBindings = new SimpleBindings();
      simpleBindings.add(s.getField(), fromSortField(s));
      boolean reverse = s.getType() == SortField.Type.SCORE || s.getReverse();
      mutated[i] = expr.getSortField(simpleBindings, reverse);
    } else {
      mutated[i] = original[i];
    }
  }
  
  Sort mutatedSort = new Sort(mutated);
  TopDocs actual = searcher.search(query, size, mutatedSort, random().nextBoolean());
  CheckHits.checkEqual(query, expected.scoreDocs, actual.scoreDocs);
  
  if (size < actual.totalHits.value) {
    expected = searcher.searchAfter(expected.scoreDocs[size-1], query, size, sort);
    actual = searcher.searchAfter(actual.scoreDocs[size-1], query, size, mutatedSort);
    CheckHits.checkEqual(query, expected.scoreDocs, actual.scoreDocs);
  }
}
 

/** Returns the grouped results.  Returns null if the
 *  number of groups collected is <= groupOffset.
 *
 *  <p><b>NOTE</b>: This collector is unable to compute
 *  the groupValue per group so it will always be null.
 *  This is normally not a problem, as you can obtain the
 *  value just like you obtain other values for each
 *  matching document (eg, via stored fields, via
 *  DocValues, etc.)
 *
 *  @param withinGroupSort The {@link Sort} used to sort
 *    documents within each group.
 *  @param groupOffset Which group to start from
 *  @param withinGroupOffset Which document to start from
 *    within each group
 *  @param maxDocsPerGroup How many top documents to keep
 *     within each group.
 */
public TopGroups<?> getTopGroups(Sort withinGroupSort, int groupOffset, int withinGroupOffset, int maxDocsPerGroup) throws IOException {

  //if (queueFull) {
  //System.out.println("getTopGroups groupOffset=" + groupOffset + " topNGroups=" + topNGroups);
  //}
  if (subDocUpto != 0) {
    processGroup();
  }
  if (groupOffset >= groupQueue.size()) {
    return null;
  }
  int totalGroupedHitCount = 0;

  final ScoreAndDoc fakeScorer = new ScoreAndDoc();

  float maxScore = Float.MIN_VALUE;

  @SuppressWarnings({"unchecked","rawtypes"})
  final GroupDocs<Object>[] groups = new GroupDocs[groupQueue.size() - groupOffset];
  for(int downTo=groupQueue.size()-groupOffset-1;downTo>=0;downTo--) {
    final OneGroup og = groupQueue.pop();

    // At this point we hold all docs w/ in each group,
    // unsorted; we now sort them:
    final TopDocsCollector<?> collector;
    if (withinGroupSort.equals(Sort.RELEVANCE)) {
      // Sort by score
      if (!needsScores) {
        throw new IllegalArgumentException("cannot sort by relevance within group: needsScores=false");
      }
      collector = TopScoreDocCollector.create(maxDocsPerGroup, Integer.MAX_VALUE);
    } else {
      // Sort by fields
      collector = TopFieldCollector.create(withinGroupSort, maxDocsPerGroup, Integer.MAX_VALUE); // TODO: disable exact counts?
    }

    float groupMaxScore = needsScores ? Float.NEGATIVE_INFINITY : Float.NaN;
    LeafCollector leafCollector = collector.getLeafCollector(og.readerContext);
    leafCollector.setScorer(fakeScorer);
    for(int docIDX=0;docIDX<og.count;docIDX++) {
      final int doc = og.docs[docIDX];
      fakeScorer.doc = doc;
      if (needsScores) {
        fakeScorer.score = og.scores[docIDX];
        groupMaxScore = Math.max(groupMaxScore, fakeScorer.score);
      }
      leafCollector.collect(doc);
    }
    totalGroupedHitCount += og.count;

    final Object[] groupSortValues;

    groupSortValues = new Comparable<?>[comparators.length];
    for(int sortFieldIDX=0;sortFieldIDX<comparators.length;sortFieldIDX++) {
      groupSortValues[sortFieldIDX] = comparators[sortFieldIDX].value(og.comparatorSlot);
    }

    final TopDocs topDocs = collector.topDocs(withinGroupOffset, maxDocsPerGroup);

    // TODO: we could aggregate scores across children
    // by Sum/Avg instead of passing NaN:
    groups[downTo] = new GroupDocs<>(Float.NaN,
                                           groupMaxScore,
                                           new TotalHits(og.count, TotalHits.Relation.EQUAL_TO),
                                           topDocs.scoreDocs,
                                           null,
                                           groupSortValues);
    maxScore = Math.max(maxScore, groupMaxScore);
  }

  /*
  while (groupQueue.size() != 0) {
    final OneGroup og = groupQueue.pop();
    //System.out.println("  leftover: og ord=" + og.groupOrd + " count=" + og.count);
    totalGroupedHitCount += og.count;
  }
  */

  return new TopGroups<>(new TopGroups<>(groupSort.getSort(),
                                     withinGroupSort.getSort(),
                                     totalHitCount, totalGroupedHitCount, groups, maxScore),
                       totalGroupCount);
}
 

/**
 * Tests index sort order.
 * @lucene.experimental
 */
public static Status.IndexSortStatus testSort(CodecReader reader, Sort sort, PrintStream infoStream, boolean failFast) throws IOException {
  // This segment claims its documents are sorted according to the incoming sort ... let's make sure:

  long startNS = System.nanoTime();

  Status.IndexSortStatus status = new Status.IndexSortStatus();

  if (sort != null) {
    if (infoStream != null) {
      infoStream.print("    test: index sort..........");
    }

    SortField fields[] = sort.getSort();
    final int reverseMul[] = new int[fields.length];
    final LeafFieldComparator comparators[] = new LeafFieldComparator[fields.length];

    LeafReaderContext readerContext = new LeafReaderContext(reader);
  
    for (int i = 0; i < fields.length; i++) {
      reverseMul[i] = fields[i].getReverse() ? -1 : 1;
      comparators[i] = fields[i].getComparator(1, i).getLeafComparator(readerContext);
    }

    int maxDoc = reader.maxDoc();

    try {

      for(int docID=1;docID < maxDoc;docID++) {
    
        int cmp = 0;

        for (int i = 0; i < comparators.length; i++) {
          // TODO: would be better if copy() didnt cause a term lookup in TermOrdVal & co,
          // the segments are always the same here...
          comparators[i].copy(0, docID-1);
          comparators[i].setBottom(0);
          cmp = reverseMul[i] * comparators[i].compareBottom(docID);
          if (cmp != 0) {
            break;
          }
        }

        if (cmp > 0) {
          throw new RuntimeException("segment has indexSort=" + sort + " but docID=" + (docID-1) + " sorts after docID=" + docID);
        }
      }
      msg(infoStream, String.format(Locale.ROOT, "OK [took %.3f sec]", nsToSec(System.nanoTime()-startNS)));
    } catch (Throwable e) {
      if (failFast) {
        throw IOUtils.rethrowAlways(e);
      }
      msg(infoStream, "ERROR [" + String.valueOf(e.getMessage()) + "]");
      status.error = e;
      if (infoStream != null) {
        e.printStackTrace(infoStream);
      }
    }
  }

  return status;
}
 

/**
 * Generates an empty CursorMark bound for use with the 
 * specified schema and {@link SortSpec}.
 *
 * @param schema used for basic validation
 * @param sortSpec bound to this totem (un)marshalling serialized values
 */
public CursorMark(IndexSchema schema, SortSpec sortSpec) {

  final SchemaField uniqueKey = schema.getUniqueKeyField();
  if (null == uniqueKey) {
    throw new SolrException(ErrorCode.BAD_REQUEST,
                            "Cursor functionality is not available unless the IndexSchema defines a uniqueKey field");
  }

  final Sort sort = sortSpec.getSort();
  if (null == sort) {
    // pure score, by definition we don't include the mandatyr uniqueKey tie breaker
    throw new SolrException(ErrorCode.BAD_REQUEST,
                            "Cursor functionality requires a sort containing a uniqueKey field tie breaker");
  }
  
  if (!sortSpec.getSchemaFields().contains(uniqueKey)) {
    throw new SolrException(ErrorCode.BAD_REQUEST,
                            "Cursor functionality requires a sort containing a uniqueKey field tie breaker");
  }

  if (0 != sortSpec.getOffset()) {
    throw new SolrException(ErrorCode.BAD_REQUEST,
                            "Cursor functionality requires start=0");
  }

  for (SortField sf : sort.getSort()) {
    if (sf.getType().equals(SortField.Type.DOC)) {
      throw new SolrException(ErrorCode.BAD_REQUEST,
                              "Cursor functionality can not be used with internal doc ordering sort: _docid_");
    }
  }

  if (sort.getSort().length != sortSpec.getSchemaFields().size()) {
      throw new SolrException(ErrorCode.SERVER_ERROR,
                              "Cursor SortSpec failure: sort length != SchemaFields: " 
                              + sort.getSort().length + " != " + 
                              sortSpec.getSchemaFields().size());
  }

  this.sortSpec = sortSpec;
  this.values = null;
}