Index: lucene/join/src/test/org/apache/lucene/search/join/TestJoinUtil.java IDEA additional info: Subsystem: com.intellij.openapi.diff.impl.patch.CharsetEP <+>UTF-8 =================================================================== --- lucene/join/src/test/org/apache/lucene/search/join/TestJoinUtil.java (revision 1335343) +++ lucene/join/src/test/org/apache/lucene/search/join/TestJoinUtil.java (revision ) @@ -95,21 +95,21 @@ // Search for product Query joinQuery = - JoinUtil.createJoinQuery(idField, false, toField, new TermQuery(new Term("name", "name2")), indexSearcher); + JoinUtil.createJoinQuery(idField, false, toField, new TermQuery(new Term("name", "name2")), indexSearcher, ToParentBlockJoinQuery.ScoreMode.None); TopDocs result = indexSearcher.search(joinQuery, 10); assertEquals(2, result.totalHits); assertEquals(4, result.scoreDocs[0].doc); assertEquals(5, result.scoreDocs[1].doc); - joinQuery = JoinUtil.createJoinQuery(idField, false, toField, new TermQuery(new Term("name", "name1")), indexSearcher); + joinQuery = JoinUtil.createJoinQuery(idField, false, toField, new TermQuery(new Term("name", "name1")), indexSearcher, ToParentBlockJoinQuery.ScoreMode.None); result = indexSearcher.search(joinQuery, 10); assertEquals(2, result.totalHits); assertEquals(1, result.scoreDocs[0].doc); assertEquals(2, result.scoreDocs[1].doc); // Search for offer - joinQuery = JoinUtil.createJoinQuery(toField, false, idField, new TermQuery(new Term("id", "5")), indexSearcher); + joinQuery = JoinUtil.createJoinQuery(toField, false, idField, new TermQuery(new Term("id", "5")), indexSearcher, ToParentBlockJoinQuery.ScoreMode.None); result = indexSearcher.search(joinQuery, 10); assertEquals(1, result.totalHits); assertEquals(3, result.scoreDocs[0].doc); @@ -118,6 +118,96 @@ dir.close(); } + public void testSimpleWithScoring() throws Exception { + final String idField = "id"; + final String toField = "movieId"; + + Directory dir = newDirectory(); + RandomIndexWriter w = new RandomIndexWriter( + random(), + dir, + newIndexWriterConfig(TEST_VERSION_CURRENT, + new MockAnalyzer(random())).setMergePolicy(newLogMergePolicy())); + + // 0 + Document doc = new Document(); + doc.add(new Field("description", "A random movie", TextField.TYPE_STORED)); + doc.add(new Field("name", "Movie 1", TextField.TYPE_STORED)); + doc.add(new Field(idField, "1", TextField.TYPE_STORED)); + w.addDocument(doc); + + // 1 + doc = new Document(); + doc.add(new Field("subtitle", "The first subtitle of this movie", TextField.TYPE_STORED)); + doc.add(new Field(idField, "2", TextField.TYPE_STORED)); + doc.add(new Field(toField, "1", TextField.TYPE_STORED)); + w.addDocument(doc); + + // 2 + doc = new Document(); + doc.add(new Field("subtitle", "random subtitle; random event movie", TextField.TYPE_STORED)); + doc.add(new Field(idField, "3", TextField.TYPE_STORED)); + doc.add(new Field(toField, "1", TextField.TYPE_STORED)); + w.addDocument(doc); + + // 3 + doc = new Document(); + doc.add(new Field("description", "A second random movie", TextField.TYPE_STORED)); + doc.add(new Field("name", "Movie 2", TextField.TYPE_STORED)); + doc.add(new Field(idField, "4", TextField.TYPE_STORED)); + w.addDocument(doc); + w.commit(); + + // 4 + doc = new Document(); + doc.add(new Field("subtitle", "a very random event happened during christmas night", TextField.TYPE_STORED)); + doc.add(new Field(idField, "5", TextField.TYPE_STORED)); + doc.add(new Field(toField, "4", TextField.TYPE_STORED)); + w.addDocument(doc); + + // 5 + doc = new Document(); + doc.add(new Field("subtitle", "movie end movie test 123 test 123 random", TextField.TYPE_STORED)); + doc.add(new Field(idField, "6", TextField.TYPE_STORED)); + doc.add(new Field(toField, "4", TextField.TYPE_STORED)); + w.addDocument(doc); + + IndexSearcher indexSearcher = new IndexSearcher(w.getReader()); + w.close(); + + // Search for movie via subtitle + Query joinQuery = + JoinUtil.createJoinQuery(toField, false, idField, new TermQuery(new Term("subtitle", "random")), indexSearcher, ToParentBlockJoinQuery.ScoreMode.Max); + TopDocs result = indexSearcher.search(joinQuery, 10); + assertEquals(2, result.totalHits); + assertEquals(0, result.scoreDocs[0].doc); + assertEquals(3, result.scoreDocs[1].doc); + + // Score mode max. + joinQuery = JoinUtil.createJoinQuery(toField, false, idField, new TermQuery(new Term("subtitle", "movie")), indexSearcher, ToParentBlockJoinQuery.ScoreMode.Max); + result = indexSearcher.search(joinQuery, 10); + assertEquals(2, result.totalHits); + assertEquals(3, result.scoreDocs[0].doc); + assertEquals(0, result.scoreDocs[1].doc); + + // Score mode total + joinQuery = JoinUtil.createJoinQuery(toField, false, idField, new TermQuery(new Term("subtitle", "movie")), indexSearcher, ToParentBlockJoinQuery.ScoreMode.Total); + result = indexSearcher.search(joinQuery, 10); + assertEquals(2, result.totalHits); + assertEquals(0, result.scoreDocs[0].doc); + assertEquals(3, result.scoreDocs[1].doc); + + //Score mode avg + joinQuery = JoinUtil.createJoinQuery(toField, false, idField, new TermQuery(new Term("subtitle", "movie")), indexSearcher, ToParentBlockJoinQuery.ScoreMode.Avg); + result = indexSearcher.search(joinQuery, 10); + assertEquals(2, result.totalHits); + assertEquals(0, result.scoreDocs[0].doc); + assertEquals(3, result.scoreDocs[1].doc); + + indexSearcher.getIndexReader().close(); + dir.close(); + } + @Test public void testSingleValueRandomJoin() throws Exception { int maxIndexIter = _TestUtil.nextInt(random(), 6, 12); @@ -166,9 +256,9 @@ } Query joinQuery; if (from) { - joinQuery = JoinUtil.createJoinQuery("from", multipleValuesPerDocument, "to", actualQuery, indexSearcher); + joinQuery = JoinUtil.createJoinQuery("from", multipleValuesPerDocument, "to", actualQuery, indexSearcher, ToParentBlockJoinQuery.ScoreMode.None); } else { - joinQuery = JoinUtil.createJoinQuery("to", multipleValuesPerDocument, "from", actualQuery, indexSearcher); + joinQuery = JoinUtil.createJoinQuery("to", multipleValuesPerDocument, "from", actualQuery, indexSearcher, ToParentBlockJoinQuery.ScoreMode.None); } if (VERBOSE) { System.out.println("joinQuery=" + joinQuery); Index: lucene/join/src/java/org/apache/lucene/search/join/TermsIncludingScoreQuery.java IDEA additional info: Subsystem: com.intellij.openapi.diff.impl.patch.CharsetEP <+>UTF-8 =================================================================== --- lucene/join/src/java/org/apache/lucene/search/join/TermsIncludingScoreQuery.java (revision ) +++ lucene/join/src/java/org/apache/lucene/search/join/TermsIncludingScoreQuery.java (revision ) @@ -0,0 +1,154 @@ +package org.apache.lucene.search.join; + +/* + * 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 org.apache.lucene.index.AtomicReaderContext; +import org.apache.lucene.index.DocsEnum; +import org.apache.lucene.index.TermsEnum; +import org.apache.lucene.search.*; +import org.apache.lucene.util.Bits; +import org.apache.lucene.util.BytesRef; +import org.apache.lucene.util.BytesRefHash; + +import java.io.IOException; + +class TermsIncludingScoreQuery extends Query { + + final String field; + final BytesRefHash terms; + final float[] scores; + final int[] ords; + + TermsIncludingScoreQuery(String field, BytesRefHash terms, float[] scores) { + this.field = field; + this.terms = terms; + this.scores = scores; + this.ords = terms.sort(BytesRef.getUTF8SortedAsUnicodeComparator()); + } + + public String toString(String string) { + return "TermsIncludingScoreQuery{" + + "field=" + field + + '}'; + } + + public static TermsIncludingScoreQuery create(String field, BytesRefHash terms, float[] scores, boolean multipleValuesPerDocument) { + if (multipleValuesPerDocument) { + throw new UnsupportedOperationException("Multiple values not yet supported"); + } else { + return new SV(field, terms, scores); + } + } + + static class SV extends TermsIncludingScoreQuery { + + SV(String field, BytesRefHash terms, float[] scores) { + super(field, terms, scores); + } + + @Override + public Weight createWeight(IndexSearcher searcher) throws IOException { + return new Weight() { + public Explanation explain(AtomicReaderContext context, int doc) throws IOException { + return null; + } + + public Query getQuery() { + return SV.this; + } + + public float getValueForNormalization() throws IOException { + return 0; + } + + public void normalize(float norm, float topLevelBoost) { + } + + public Scorer scorer(AtomicReaderContext context, boolean scoreDocsInOrder, boolean topScorer, Bits acceptDocs) throws IOException { + return new SVInnerScorer(this, acceptDocs, context.reader().terms(field).iterator(null)); + } + }; + } + + } + + // This impl assumes that the 'join' values are used uniquely per doc per field. + // + class SVInnerScorer extends Scorer { + + private final BytesRef spare = new BytesRef(); + private final Bits acceptDocs; + private final TermsEnum termsEnum; + + private DocsEnum docsEnum; + + private int upto; + private int currentUptoForScore; + + SVInnerScorer(Weight weight, Bits acceptDocs, TermsEnum termsEnum) { + super(weight); + this.acceptDocs = acceptDocs; + this.termsEnum = termsEnum; + } + + public float score() throws IOException { + return scores[ords[currentUptoForScore]]; + } + + public int docID() { + return docsEnum.docID(); + } + + public int nextDoc() throws IOException { + if (docsEnum != null) { + int docId = docsEnum.nextDoc(); + if (docId == DocIdSetIterator.NO_MORE_DOCS) { + docsEnum = null; + } else { + return docId; + } + } + + do { + if (upto == terms.size()) { + return DocIdSetIterator.NO_MORE_DOCS; + } + + currentUptoForScore = upto; + TermsEnum.SeekStatus status = termsEnum.seekCeil(terms.get(ords[upto++], spare), true); + if (status == TermsEnum.SeekStatus.FOUND) { + docsEnum = termsEnum.docs(acceptDocs, null, false); + } + } while (docsEnum == null); + + return docsEnum.nextDoc(); + } + + public int advance(int target) throws IOException { + if (docsEnum == null) { + int docId = nextDoc(); + if (docId >= target) { + return docId; + } + } + + return docsEnum.advance(target); + } + } + +} Index: lucene/join/src/java/org/apache/lucene/search/join/TermsWithScoreCollector.java IDEA additional info: Subsystem: com.intellij.openapi.diff.impl.patch.CharsetEP <+>UTF-8 =================================================================== --- lucene/join/src/java/org/apache/lucene/search/join/TermsWithScoreCollector.java (revision ) +++ lucene/join/src/java/org/apache/lucene/search/join/TermsWithScoreCollector.java (revision ) @@ -0,0 +1,160 @@ +package org.apache.lucene.search.join; + +/* + * 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 org.apache.lucene.index.AtomicReaderContext; +import org.apache.lucene.search.Collector; +import org.apache.lucene.search.FieldCache; +import org.apache.lucene.search.Scorer; +import org.apache.lucene.util.ArrayUtil; +import org.apache.lucene.util.BytesRef; +import org.apache.lucene.util.BytesRefHash; + +import java.io.IOException; + +abstract class TermsWithScoreCollector extends Collector { + + private final static int INITIAL_GROW_FACTOR = 2; + + final String field; + final BytesRefHash collectorTerms = new BytesRefHash(); + final ToParentBlockJoinQuery.ScoreMode scoreMode; + + Scorer scorer; + int growFactor = INITIAL_GROW_FACTOR; + float[] scores = new float[256]; + + TermsWithScoreCollector(String field, ToParentBlockJoinQuery.ScoreMode scoreMode) { + this.field = field; + this.scoreMode = scoreMode; + } + + public BytesRefHash getCollectorTerms() { + return collectorTerms; + } + + public float[] getScoresPerTerm() { + return scores; + } + + public void setScorer(Scorer scorer) throws IOException { + this.scorer = scorer; + } + + public boolean acceptsDocsOutOfOrder() { + return true; + } + + /** + * Chooses the right {@link TermsWithScoreCollector} implementation. + * + * @param field The field to collect terms for + * @param multipleValuesPerDocument Whether the field to collect terms for has multiple values per document. + * @return a {@link TermsWithScoreCollector} instance + */ + static TermsWithScoreCollector create(String field, boolean multipleValuesPerDocument, ToParentBlockJoinQuery.ScoreMode scoreMode) { + if (multipleValuesPerDocument) { + throw new UnsupportedOperationException("Multiple values per doc not yet supported"); + } else { + switch (scoreMode) { + case Avg: + return new SVAvg(field); + default: + return new SV(field, scoreMode); + } + } + } + + // impl that works with single value per document + static class SV extends TermsWithScoreCollector { + + final BytesRef spare = new BytesRef(); + FieldCache.DocTerms fromDocTerms; + + SV(String field, ToParentBlockJoinQuery.ScoreMode scoreMode) { + super(field, scoreMode); + } + + public void collect(int doc) throws IOException { + int ord = collectorTerms.add(fromDocTerms.getTerm(doc, spare)); + if (ord < 0) { + ord = -ord - 1; + } else { + if (ord >= scores.length) { + scores = ArrayUtil.grow(scores, scores.length * growFactor); + growFactor++; + } + } + + float current = scorer.score(); + float existing = scores[ord]; + if (Float.compare(existing, 0.0f) == 0) { + scores[ord] = current; + } else { + switch (scoreMode) { + case Total: + scores[ord] = scores[ord] + current; + break; + case Max: + if (current > existing) { + scores[ord] = current; + } + } + } + } + + public void setNextReader(AtomicReaderContext context) throws IOException { + fromDocTerms = FieldCache.DEFAULT.getTerms(context.reader(), field); + } + } + + static class SVAvg extends SV { + + int[] ordScores = new int[256]; + + SVAvg(String field) { + super(field, ToParentBlockJoinQuery.ScoreMode.Avg); + } + + @Override + public void collect(int doc) throws IOException { + int ord = collectorTerms.add(fromDocTerms.getTerm(doc, spare)); + if (ord < 0) { + ord = -ord - 1; + } else { + if (ord >= scores.length) { + scores = ArrayUtil.grow(scores, scores.length * growFactor); + ordScores = ArrayUtil.grow(ordScores, ordScores.length * growFactor); + growFactor++; + } + } + + float current = scorer.score(); + float existing = scores[ord]; + if (Float.compare(existing, 0.0f) == 0) { + scores[ord] = current; + ordScores[ord] = 1; + } else { + scores[ord] = scores[ord] + current; + ordScores[ord]++; + } + } + + } + +} Index: lucene/join/src/java/org/apache/lucene/search/join/JoinUtil.java IDEA additional info: Subsystem: com.intellij.openapi.diff.impl.patch.CharsetEP <+>UTF-8 =================================================================== --- lucene/join/src/java/org/apache/lucene/search/join/JoinUtil.java (revision 1335343) +++ lucene/join/src/java/org/apache/lucene/search/join/JoinUtil.java (revision ) @@ -39,11 +39,13 @@ * Execute the returned query with a {@link IndexSearcher} to retrieve all documents that have the same terms in the * to field that match with documents matching the specified fromQuery and have the same terms in the from field. * + * * @param fromField The from field to join from * @param multipleValuesPerDocument Whether the from field has multiple terms per document * @param toField The to field to join to * @param fromQuery The query to match documents on the from side * @param fromSearcher The searcher that executed the specified fromQuery + * @param scoreMode * @return a {@link Query} instance that can be used to join documents based on the * terms in the from and to field * @throws IOException If I/O related errors occur @@ -52,10 +54,18 @@ boolean multipleValuesPerDocument, String toField, Query fromQuery, - IndexSearcher fromSearcher) throws IOException { + IndexSearcher fromSearcher, + ToParentBlockJoinQuery.ScoreMode scoreMode) throws IOException { + switch (scoreMode) { + case None: - TermsCollector termsCollector = TermsCollector.create(fromField, multipleValuesPerDocument); - fromSearcher.search(fromQuery, termsCollector); - return new TermsQuery(toField, termsCollector.getCollectorTerms()); + TermsCollector termsCollector = TermsCollector.create(fromField, multipleValuesPerDocument); + fromSearcher.search(fromQuery, termsCollector); + return new TermsQuery(toField, termsCollector.getCollectorTerms()); + default: + TermsWithScoreCollector termsWithScoreCollector = TermsWithScoreCollector.create(fromField, multipleValuesPerDocument, scoreMode); + fromSearcher.search(fromQuery, termsWithScoreCollector); + return TermsIncludingScoreQuery.create(toField, termsWithScoreCollector.getCollectorTerms(), termsWithScoreCollector.getScoresPerTerm(), multipleValuesPerDocument); + } } }