--- /dev/null
+package org.apache.lucene.search.grouping;
+
+/**
+ * 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 java.io.IOException;
+
+import org.apache.lucene.index.IndexReader;
+import org.apache.lucene.index.IndexWriter; // javadocs
+import org.apache.lucene.search.Collector;
+import org.apache.lucene.search.DocIdSetIterator;
+import org.apache.lucene.search.FieldComparator;
+import org.apache.lucene.search.Filter;
+import org.apache.lucene.search.Scorer;
+import org.apache.lucene.search.Sort;
+import org.apache.lucene.search.SortField;
+import org.apache.lucene.search.TopDocs;
+import org.apache.lucene.search.TopDocsCollector;
+import org.apache.lucene.search.TopFieldCollector;
+import org.apache.lucene.search.TopScoreDocCollector;
+import org.apache.lucene.search.Weight;
+import org.apache.lucene.util.ArrayUtil;
+import org.apache.lucene.util.PriorityQueue;
+
+/** BlockGroupingCollector performs grouping with a
+ * single pass collector, as long as you are grouping by a
+ * doc block field, ie all documents sharing a given group
+ * value were indexed as a doc block using the atomic
+ * {@link IndexWriter#addDocuments} or {@link
+ * IndexWriter#updateDocuments} API.
+ *
+ * <p>This results in faster performance (~25% faster QPS)
+ * than the two-pass grouping collectors, with the tradeoff
+ * being that the documents in each group must always be
+ * indexed as a block. This collector also fills in
+ * TopGroups.totalGroupCount without requiring the separate
+ * {@link TermAllGroupsCollector}. However, this collector does
+ * not fill in the groupValue of each group; this field
+ * will always be null.
+ *
+ * <p><b>NOTE</b>: this collector makes no effort to verify
+ * the docs were in fact indexed as a block, so it's up to
+ * you to ensure this was the case.
+ *
+ * <p>See {@link org.apache.lucene.search.grouping} for more
+ * details including a full code example.</p>
+ *
+ * @lucene.experimental
+ */
+
+public class BlockGroupingCollector extends Collector {
+
+ private int[] pendingSubDocs;
+ private float[] pendingSubScores;
+ private int subDocUpto;
+
+ private final Sort groupSort;
+ private final int topNGroups;
+ private final Filter lastDocPerGroup;
+
+ // TODO: specialize into 2 classes, static "create" method:
+ private final boolean needsScores;
+
+ private final FieldComparator[] comparators;
+ private final int[] reversed;
+ private final int compIDXEnd;
+ private int bottomSlot;
+ private boolean queueFull;
+ private IndexReader currentReader;
+
+ private int topGroupDoc;
+ private int totalHitCount;
+ private int totalGroupCount;
+ private int docBase;
+ private int groupEndDocID;
+ private DocIdSetIterator lastDocPerGroupBits;
+ private Scorer scorer;
+ private final GroupQueue groupQueue;
+ private boolean groupCompetes;
+
+ private final static class FakeScorer extends Scorer {
+
+ float score;
+ int doc;
+
+ public FakeScorer() {
+ super((Weight) null);
+ }
+
+ @Override
+ public float score() {
+ return score;
+ }
+
+ @Override
+ public int docID() {
+ return doc;
+ }
+
+ @Override
+ public int advance(int target) {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public int nextDoc() {
+ throw new UnsupportedOperationException();
+ }
+ }
+
+ private static final class OneGroup {
+ IndexReader reader;
+ int docBase;
+ //int groupOrd;
+ int topGroupDoc;
+ int[] docs;
+ float[] scores;
+ int count;
+ int comparatorSlot;
+ }
+
+ // Sorts by groupSort. Not static -- uses comparators, reversed
+ private final class GroupQueue extends PriorityQueue<OneGroup> {
+
+ public GroupQueue(int size) {
+ initialize(size);
+ }
+
+ @Override
+ protected boolean lessThan(final OneGroup group1, final OneGroup group2) {
+
+ //System.out.println(" ltcheck");
+ assert group1 != group2;
+ assert group1.comparatorSlot != group2.comparatorSlot;
+
+ final int numComparators = comparators.length;
+ for (int compIDX=0;compIDX < numComparators; compIDX++) {
+ final int c = reversed[compIDX] * comparators[compIDX].compare(group1.comparatorSlot, group2.comparatorSlot);
+ if (c != 0) {
+ // Short circuit
+ return c > 0;
+ }
+ }
+
+ // Break ties by docID; lower docID is always sorted first
+ return group1.topGroupDoc > group2.topGroupDoc;
+ }
+ }
+
+ // Called when we transition to another group; if the
+ // group is competitive we insert into the group queue
+ private void processGroup() {
+ totalGroupCount++;
+ //System.out.println(" processGroup ord=" + lastGroupOrd + " competes=" + groupCompetes + " count=" + subDocUpto + " groupDoc=" + topGroupDoc);
+ if (groupCompetes) {
+ if (!queueFull) {
+ // Startup transient: always add a new OneGroup
+ final OneGroup og = new OneGroup();
+ og.count = subDocUpto;
+ og.topGroupDoc = docBase + topGroupDoc;
+ og.docs = pendingSubDocs;
+ pendingSubDocs = new int[10];
+ if (needsScores) {
+ og.scores = pendingSubScores;
+ pendingSubScores = new float[10];
+ }
+ og.reader = currentReader;
+ og.docBase = docBase;
+ //og.groupOrd = lastGroupOrd;
+ og.comparatorSlot = bottomSlot;
+ final OneGroup bottomGroup = groupQueue.add(og);
+ //System.out.println(" ADD group=" + getGroupString(lastGroupOrd) + " newBottom=" + getGroupString(bottomGroup.groupOrd));
+ queueFull = groupQueue.size() == topNGroups;
+ if (queueFull) {
+ // Queue just became full; now set the real bottom
+ // in the comparators:
+ bottomSlot = bottomGroup.comparatorSlot;
+ //System.out.println(" set bottom=" + bottomSlot);
+ for (int i = 0; i < comparators.length; i++) {
+ comparators[i].setBottom(bottomSlot);
+ }
+ //System.out.println(" QUEUE FULL");
+ } else {
+ // Queue not full yet -- just advance bottomSlot:
+ bottomSlot = groupQueue.size();
+ }
+ } else {
+ // Replace bottom element in PQ and then updateTop
+ final OneGroup og = groupQueue.top();
+ assert og != null;
+ og.count = subDocUpto;
+ og.topGroupDoc = docBase + topGroupDoc;
+ // Swap pending docs
+ final int[] savDocs = og.docs;
+ og.docs = pendingSubDocs;
+ pendingSubDocs = savDocs;
+ if (needsScores) {
+ // Swap pending scores
+ final float[] savScores = og.scores;
+ og.scores = pendingSubScores;
+ pendingSubScores = savScores;
+ }
+ og.reader = currentReader;
+ og.docBase = docBase;
+ //og.groupOrd = lastGroupOrd;
+ bottomSlot = groupQueue.updateTop().comparatorSlot;
+
+ //System.out.println(" set bottom=" + bottomSlot);
+ for (int i = 0; i < comparators.length; i++) {
+ comparators[i].setBottom(bottomSlot);
+ }
+ }
+ }
+ subDocUpto = 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 Filter} that marks the
+ * last document in each group.
+ */
+ public BlockGroupingCollector(Sort groupSort, int topNGroups, boolean needsScores, Filter lastDocPerGroup) throws IOException {
+
+ 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;
+ // TODO: allow null groupSort to mean "by relevance",
+ // and specialize it?
+ this.groupSort = groupSort;
+
+ this.topNGroups = topNGroups;
+
+ final SortField[] sortFields = groupSort.getSort();
+ comparators = new FieldComparator[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;
+ }
+ }
+
+ // TODO: maybe allow no sort on retrieving groups? app
+ // may want to simply process docs in the group itself?
+ // typically they will be presented as a "single" result
+ // in the UI?
+
+ /** 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
+ * FieldCache, etc.)
+ *
+ * @param withinGroupSort The {@link Sort} used to sort
+ * documents within each group. Passing null is
+ * allowed, to sort by relevance.
+ * @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.
+ * @param fillSortFields If true then the Comparable
+ * values for the sort fields will be set
+ */
+ public TopGroups getTopGroups(Sort withinGroupSort, int groupOffset, int withinGroupOffset, int maxDocsPerGroup, boolean fillSortFields) 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 FakeScorer fakeScorer = new FakeScorer();
+
+ @SuppressWarnings("unchecked")
+ 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 == null) {
+ // Sort by score
+ if (!needsScores) {
+ throw new IllegalArgumentException("cannot sort by relevance within group: needsScores=false");
+ }
+ collector = TopScoreDocCollector.create(maxDocsPerGroup, true);
+ } else {
+ // Sort by fields
+ collector = TopFieldCollector.create(withinGroupSort, maxDocsPerGroup, fillSortFields, needsScores, needsScores, true);
+ }
+
+ collector.setScorer(fakeScorer);
+ collector.setNextReader(og.reader, og.docBase);
+ for(int docIDX=0;docIDX<og.count;docIDX++) {
+ final int doc = og.docs[docIDX];
+ fakeScorer.doc = doc;
+ if (needsScores) {
+ fakeScorer.score = og.scores[docIDX];
+ }
+ collector.collect(doc);
+ }
+ totalGroupedHitCount += og.count;
+
+ final Object[] groupSortValues;
+
+ if (fillSortFields) {
+ groupSortValues = new Comparable[comparators.length];
+ for(int sortFieldIDX=0;sortFieldIDX<comparators.length;sortFieldIDX++) {
+ groupSortValues[sortFieldIDX] = comparators[sortFieldIDX].value(og.comparatorSlot);
+ }
+ } else {
+ groupSortValues = null;
+ }
+
+ final TopDocs topDocs = collector.topDocs(withinGroupOffset, maxDocsPerGroup);
+
+ groups[downTo] = new GroupDocs<Object>(topDocs.getMaxScore(),
+ og.count,
+ topDocs.scoreDocs,
+ null,
+ groupSortValues);
+ }
+
+ /*
+ 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<Object>(new TopGroups<Object>(groupSort.getSort(),
+ withinGroupSort == null ? null : withinGroupSort.getSort(),
+ totalHitCount, totalGroupedHitCount, groups),
+ totalGroupCount);
+ }
+
+ @Override
+ public void setScorer(Scorer scorer) throws IOException {
+ this.scorer = scorer;
+ for (FieldComparator comparator : comparators) {
+ comparator.setScorer(scorer);
+ }
+ }
+
+ @Override
+ public void collect(int doc) throws IOException {
+
+ // System.out.println("C " + doc);
+
+ if (doc > groupEndDocID) {
+ // Group changed
+ if (subDocUpto != 0) {
+ processGroup();
+ }
+ groupEndDocID = lastDocPerGroupBits.advance(doc);
+ //System.out.println(" adv " + groupEndDocID + " " + lastDocPerGroupBits);
+ subDocUpto = 0;
+ groupCompetes = !queueFull;
+ }
+
+ totalHitCount++;
+
+ // Always cache doc/score within this group:
+ if (subDocUpto == pendingSubDocs.length) {
+ pendingSubDocs = ArrayUtil.grow(pendingSubDocs);
+ }
+ pendingSubDocs[subDocUpto] = doc;
+ if (needsScores) {
+ if (subDocUpto == pendingSubScores.length) {
+ pendingSubScores = ArrayUtil.grow(pendingSubScores);
+ }
+ pendingSubScores[subDocUpto] = scorer.score();
+ }
+ subDocUpto++;
+
+ if (groupCompetes) {
+ if (subDocUpto == 1) {
+ assert !queueFull;
+
+ //System.out.println(" init copy to bottomSlot=" + bottomSlot);
+ for (FieldComparator fc : comparators) {
+ fc.copy(bottomSlot, doc);
+ fc.setBottom(bottomSlot);
+ }
+ topGroupDoc = doc;
+ } else {
+ // Compare to bottomSlot
+ for (int compIDX = 0;; compIDX++) {
+ final int c = reversed[compIDX] * comparators[compIDX].compareBottom(doc);
+ if (c < 0) {
+ // Definitely not competitive -- done
+ return;
+ } else if (c > 0) {
+ // Definitely competitive.
+ break;
+ } else if (compIDX == compIDXEnd) {
+ // Ties with bottom, except we know this docID is
+ // > docID in the queue (docs are visited in
+ // order), so not competitive:
+ return;
+ }
+ }
+
+ //System.out.println(" best w/in group!");
+
+ for (FieldComparator fc : comparators) {
+ fc.copy(bottomSlot, doc);
+ // Necessary because some comparators cache
+ // details of bottom slot; this forces them to
+ // re-cache:
+ fc.setBottom(bottomSlot);
+ }
+ topGroupDoc = doc;
+ }
+ } else {
+ // We're not sure this group will make it into the
+ // queue yet
+ for (int compIDX = 0;; compIDX++) {
+ final int c = reversed[compIDX] * comparators[compIDX].compareBottom(doc);
+ if (c < 0) {
+ // Definitely not competitive -- done
+ //System.out.println(" doc doesn't compete w/ top groups");
+ return;
+ } else if (c > 0) {
+ // Definitely competitive.
+ break;
+ } else if (compIDX == compIDXEnd) {
+ // Ties with bottom, except we know this docID is
+ // > docID in the queue (docs are visited in
+ // order), so not competitive:
+ //System.out.println(" doc doesn't compete w/ top groups");
+ return;
+ }
+ }
+ groupCompetes = true;
+ for (FieldComparator fc : comparators) {
+ fc.copy(bottomSlot, doc);
+ // Necessary because some comparators cache
+ // details of bottom slot; this forces them to
+ // re-cache:
+ fc.setBottom(bottomSlot);
+ }
+ topGroupDoc = doc;
+ //System.out.println(" doc competes w/ top groups");
+ }
+ }
+
+ @Override
+ public boolean acceptsDocsOutOfOrder() {
+ return false;
+ }
+
+ @Override
+ public void setNextReader(IndexReader reader, int docBase) throws IOException {
+ if (subDocUpto != 0) {
+ processGroup();
+ }
+ subDocUpto = 0;
+ this.docBase = docBase;
+ //System.out.println("setNextReader base=" + docBase + " r=" + readerContext.reader);
+ lastDocPerGroupBits = lastDocPerGroup.getDocIdSet(reader).iterator();
+ groupEndDocID = -1;
+
+ currentReader = reader;
+ for (int i=0; i<comparators.length; i++) {
+ comparators[i].setNextReader(reader, docBase);
+ }
+ }
+}
fnp / pylucene.git / blobdiff