diff --git a/lucene/join/src/java/org/apache/lucene/search/join/ToParentBlockJoinCollector.java b/lucene/join/src/java/org/apache/lucene/search/join/ToParentBlockJoinCollector.java
index 583eee5..680f7ec 100644
--- a/lucene/join/src/java/org/apache/lucene/search/join/ToParentBlockJoinCollector.java
+++ b/lucene/join/src/java/org/apache/lucene/search/join/ToParentBlockJoinCollector.java
@@ -360,13 +360,23 @@ public class ToParentBlockJoinCollector extends Collector {
/** Return the TopGroups for the specified
* BlockJoinQuery. The groupValue of each GroupDocs will
- * be the parent docID for that group. Note that the
- * {@link GroupDocs#totalHits}, which would be the
- * total number of child documents matching that parent,
- * is not computed (will always be 0). Returns null if
- * no groups matched. */
+ * be the parent docID for that group.
+ * The number of documents within each group is bounded by maxDocsPerGroup
+ * in case of maxDocsPerGroup is not negative; otherwise all child documents will be collected.
+ * Returns null if no groups matched.
+ *
+ * @param query Search query
+ * @param withinGroupSort Sort criteria within groups
+ * @param offset Parent docs offset
+ * @param maxDocsPerGroup Upper bound of documents per group number. If it's negative all child docs will be collected
+ * @param withinGroupOffset Offset within each group of child docs
+ * @param fillSortFields Specifies whether to add sort fields or not
+ * @return TopGroups for specified query
+ * @throws IOException if there is a low-level I/O error
+ */
@SuppressWarnings("unchecked")
- public TopGroups getTopGroups(ToParentBlockJoinQuery query, Sort withinGroupSort, int offset, int maxDocsPerGroup, int withinGroupOffset, boolean fillSortFields)
+ public TopGroups getTopGroups(ToParentBlockJoinQuery query, Sort withinGroupSort, int offset,
+ int maxDocsPerGroup, int withinGroupOffset, boolean fillSortFields)
throws IOException {
@@ -391,68 +401,116 @@ public class ToParentBlockJoinCollector extends Collector {
return null;
}
- int totalGroupedHitCount = 0;
+ final GroupDocsAccumulator groupDocsAccumulator = new GroupDocsAccumulator(slot, offset, maxDocsPerGroup, withinGroupOffset,
+ withinGroupSort, fillSortFields);
+
+ groupDocsAccumulator.accumulate();
+
+ return new TopGroups(new TopGroups(sort.getSort(),
+ withinGroupSort == null ? null : withinGroupSort.getSort(),
+ 0, groupDocsAccumulator.totalGroupedHitCount, groupDocsAccumulator.groups, maxScore),
+ totalHitCount);
+ }
+
+ private final class GroupDocsAccumulator {
+ final int slot;
+ final int offset;
+ final int maxDocsPerGroup;
+ final int withinGroupOffset;
+ final Sort withinGroupSort;
+ final boolean fillSortFields;
+ final FakeScorer fakeScorer;
+ final GroupDocs[] groups;
- final FakeScorer fakeScorer = new FakeScorer();
+ int totalGroupedHitCount = 0;
@SuppressWarnings({"unchecked","rawtypes"})
- final GroupDocs[] groups = new GroupDocs[sortedGroups.length - offset];
-
- for(int groupIDX=offset;groupIDX collector;
- if (withinGroupSort == null) {
- // Sort by score
- if (!trackScores) {
- throw new IllegalArgumentException("cannot sort by relevance within group: trackScores=false");
+ public GroupDocsAccumulator(int slot, int offset, int maxDocsPerGroup, int withinGroupOffset, Sort withinGroupSort, boolean fillSortFields) {
+ this.slot = slot;
+ this.offset = offset;
+ this.maxDocsPerGroup = maxDocsPerGroup;
+ this.withinGroupOffset = withinGroupOffset;
+ this.withinGroupSort = withinGroupSort;
+ this.fillSortFields = fillSortFields;
+ this.fakeScorer = new FakeScorer();
+ this.groups = new GroupDocs[sortedGroups.length - offset];
+ }
+
+ public void accumulate() throws IOException {
+ for(int groupIDX=offset;groupIDX collector;
+ if (withinGroupSort == null) {
+ // Sort by score
+ if (!trackScores) {
+ throw new IllegalArgumentException("cannot sort by relevance within group: trackScores=false");
+ }
+ collector = TopScoreDocCollector.create(numDocsInGroup, true);
+ } else {
+ // Sort by fields
+ collector = TopFieldCollector.create(withinGroupSort, numDocsInGroup, fillSortFields, trackScores, trackMaxScore, true);
}
- collector = TopScoreDocCollector.create(maxDocsPerGroup, true);
- } else {
- // Sort by fields
- collector = TopFieldCollector.create(withinGroupSort, maxDocsPerGroup, fillSortFields, trackScores, trackMaxScore, true);
- }
- collector.setScorer(fakeScorer);
- collector.setNextReader(og.readerContext);
- final int numChildDocs = og.counts[slot];
- for(int docIDX=0;docIDX(og.score,
- topDocs.getMaxScore(),
- og.counts[slot],
- topDocs.scoreDocs,
- og.doc,
- groupSortValues);
+ groups[groupIDX-offset] = new GroupDocs(og.score,
+ topDocs.getMaxScore(),
+ numChildDocs,
+ topDocs.scoreDocs,
+ og.doc,
+ groupSortValues);
+ }
}
+ }
- return new TopGroups(new TopGroups(sort.getSort(),
- withinGroupSort == null ? null : withinGroupSort.getSort(),
- 0, totalGroupedHitCount, groups, maxScore),
- totalHitCount);
+ /** Return the TopGroups for the specified BlockJoinQuery.
+ * The groupValue of each GroupDocs will be the parent docID for that group.
+ * The number of documents within each group
+ * equals to the total number of child documents matching parent for that group.
+ * Returns null if no groups matched.
+ *
+ * @param query Search query
+ * @param withinGroupSort Sort criteria within groups
+ * @param offset Parent docs offset
+ * @param withinGroupOffset Offset within each group of child docs
+ * @param fillSortFields Specifies whether to add sort fields or not
+ * @return TopGroups for specified query
+ * @throws IOException if there is a low-level I/O error
+ */
+ @SuppressWarnings("unchecked")
+ public TopGroups getTopGroupsWithAllChildDocs(ToParentBlockJoinQuery query, Sort withinGroupSort, int offset,
+ int withinGroupOffset, boolean fillSortFields)
+ throws IOException {
+
+ return getTopGroups(query, withinGroupSort, offset, -1, withinGroupOffset, fillSortFields);
}
/**
diff --git a/lucene/join/src/test/org/apache/lucene/search/join/TestBlockJoin.java b/lucene/join/src/test/org/apache/lucene/search/join/TestBlockJoin.java
index b93a9e2..881f77c 100644
--- a/lucene/join/src/test/org/apache/lucene/search/join/TestBlockJoin.java
+++ b/lucene/join/src/test/org/apache/lucene/search/join/TestBlockJoin.java
@@ -1254,4 +1254,95 @@ public class TestBlockJoin extends LuceneTestCase {
r.close();
dir.close();
}
+
+ public void testGetTopGroups() throws Exception {
+
+ final Directory dir = newDirectory();
+ final RandomIndexWriter w = new RandomIndexWriter(random(), dir);
+
+ final List docs = new ArrayList();
+ docs.add(makeJob("ruby", 2005));
+ docs.add(makeJob("java", 2006));
+ docs.add(makeJob("java", 2010));
+ docs.add(makeJob("java", 2012));
+ Collections.shuffle(docs, random());
+ docs.add(makeResume("Frank", "United States"));
+
+ addSkillless(w);
+ w.addDocuments(docs);
+ addSkillless(w);
+
+ IndexReader r = w.getReader();
+ w.close();
+ IndexSearcher s = newSearcher(r);
+
+ // Create a filter that defines "parent" documents in the index - in this case resumes
+ Filter parentsFilter = new CachingWrapperFilter(new QueryWrapperFilter(new TermQuery(new Term("docType", "resume"))));
+
+ // Define child document criteria (finds an example of relevant work experience)
+ BooleanQuery childQuery = new BooleanQuery();
+ childQuery.add(new BooleanClause(new TermQuery(new Term("skill", "java")), Occur.MUST));
+ childQuery.add(new BooleanClause(NumericRangeQuery.newIntRange("year", 2006, 2011, true, true), Occur.MUST));
+
+ // Wrap the child document query to 'join' any matches
+ // up to corresponding parent:
+ ToParentBlockJoinQuery childJoinQuery = new ToParentBlockJoinQuery(childQuery, parentsFilter, ScoreMode.Avg);
+
+ ToParentBlockJoinCollector c = new ToParentBlockJoinCollector(Sort.RELEVANCE, 2, true, true);
+
+ s.search(childJoinQuery, c);
+
+ //Get all child documents within groups
+ @SuppressWarnings({"unchecked","rawtypes"})
+ TopGroups[] getTopGroupsResults = new TopGroups[2];
+ getTopGroupsResults[0] = c.getTopGroups(childJoinQuery, null, 0, 10, 0, true);
+ getTopGroupsResults[1] = c.getTopGroupsWithAllChildDocs(childJoinQuery, null, 0, 0, true);
+
+ for (TopGroups results : getTopGroupsResults) {
+ assertFalse(Float.isNaN(results.maxScore));
+ assertEquals(2, results.totalGroupedHitCount);
+ assertEquals(1, results.groups.length);
+
+ final GroupDocs group = results.groups[0];
+ assertEquals(2, group.totalHits);
+ assertFalse(Float.isNaN(group.score));
+ assertNotNull(group.groupValue);
+ StoredDocument parentDoc = s.doc(group.groupValue);
+ assertEquals("Frank", parentDoc.get("name"));
+
+ assertEquals(2, group.scoreDocs.length); //all matched child documents collected
+
+ for (ScoreDoc scoreDoc : group.scoreDocs) {
+ StoredDocument childDoc = s.doc(scoreDoc.doc);
+ assertEquals("java", childDoc.get("skill"));
+ int year = Integer.parseInt(childDoc.get("year"));
+ assertTrue(year >= 2006 && year <= 2011);
+ }
+ }
+
+ //Get part of child documents
+ TopGroups boundedResults = c.getTopGroups(childJoinQuery, null, 0, 1, 0, true);
+ assertFalse(Float.isNaN(boundedResults.maxScore));
+ assertEquals(2, boundedResults.totalGroupedHitCount);
+ assertEquals(1, boundedResults.groups.length);
+
+ final GroupDocs group = boundedResults.groups[0];
+ assertEquals(2, group.totalHits);
+ assertFalse(Float.isNaN(group.score));
+ assertNotNull(group.groupValue);
+ StoredDocument parentDoc = s.doc(group.groupValue);
+ assertEquals("Frank", parentDoc.get("name"));
+
+ assertEquals(1, group.scoreDocs.length); //not all matched child documents collected
+
+ for (ScoreDoc scoreDoc : group.scoreDocs) {
+ StoredDocument childDoc = s.doc(scoreDoc.doc);
+ assertEquals("java", childDoc.get("skill"));
+ int year = Integer.parseInt(childDoc.get("year"));
+ assertTrue(year >= 2006 && year <= 2011);
+ }
+
+ r.close();
+ dir.close();
+ }
}