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..d18657e 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 @@ -358,16 +358,24 @@ 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. */ - @SuppressWarnings("unchecked") - public TopGroups getTopGroups(ToParentBlockJoinQuery query, Sort withinGroupSort, int offset, int maxDocsPerGroup, int withinGroupOffset, boolean fillSortFields) - + /** Returns 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 is calculated as minimum of maxDocsPerGroup + * and number of matched child documents 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 maxDocsPerGroup Upper bound of documents per group number + * @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 + */ + public TopGroups getTopGroups(ToParentBlockJoinQuery query, Sort withinGroupSort, int offset, + int maxDocsPerGroup, int withinGroupOffset, boolean fillSortFields) throws IOException { final Integer _slot = joinQueryID.get(query); @@ -379,9 +387,6 @@ public class ToParentBlockJoinCollector extends Collector { } } - // unbox once - final int slot = _slot; - if (sortedGroups == null) { if (offset >= queue.size()) { return null; @@ -391,15 +396,43 @@ public class ToParentBlockJoinCollector extends Collector { return null; } + + final GroupDocs[] groups = accumulateGroups(_slot, offset, maxDocsPerGroup, + withinGroupOffset, withinGroupSort, fillSortFields); + int totalGroupedHitCount = 0; + for (GroupDocs group : groups) { + totalGroupedHitCount += group.totalHits; + } - final FakeScorer fakeScorer = new FakeScorer(); + return new TopGroups(new TopGroups(sort.getSort(), + withinGroupSort == null ? null : withinGroupSort.getSort(), + 0, totalGroupedHitCount, groups, maxScore), + totalHitCount); + } - @SuppressWarnings({"unchecked","rawtypes"}) + /** + * Accumulates groups for the BlockJoinQuery specified by its slot. + * + * @param slot Search query's slot + * @param offset Parent docs offset + * @param maxDocsPerGroup Upper bound of documents per group number + * @param withinGroupOffset Offset within each group of child docs + * @param withinGroupSort Sort criteria within groups + * @param fillSortFields Specifies whether to add sort fields or not + * @return groups of child documents matched by query + * @throws IOException if there is a low-level I/O error + */ + @SuppressWarnings({"unchecked","rawtypes"}) + private GroupDocs[] accumulateGroups(int slot, int offset, int maxDocsPerGroup, + int withinGroupOffset, Sort withinGroupSort, boolean fillSortFields) throws IOException { final GroupDocs[] groups = new GroupDocs[sortedGroups.length - offset]; + final FakeScorer fakeScorer = new FakeScorer(); for(int groupIDX=offset;groupIDX(og.score, topDocs.getMaxScore(), - og.counts[slot], + numChildDocs, topDocs.scoreDocs, og.doc, groupSortValues); } - return new TopGroups(new TopGroups(sort.getSort(), - withinGroupSort == null ? null : withinGroupSort.getSort(), - 0, totalGroupedHitCount, groups, maxScore), - totalHitCount); + return groups; + } + + /** Returns 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 matched child documents 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 + */ + public TopGroups getTopGroupsWithAllChildDocs(ToParentBlockJoinQuery query, Sort withinGroupSort, int offset, + int withinGroupOffset, boolean fillSortFields) + throws IOException { + + return getTopGroups(query, withinGroupSort, offset, Integer.MAX_VALUE, 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(); + } }