Index: modules/join/src/test/org/apache/lucene/search/join/TestJoinUtil.java IDEA additional info: Subsystem: com.intellij.openapi.diff.impl.patch.CharsetEP <+>UTF-8 =================================================================== --- modules/join/src/test/org/apache/lucene/search/join/TestJoinUtil.java (revision ) +++ modules/join/src/test/org/apache/lucene/search/join/TestJoinUtil.java (revision ) @@ -0,0 +1,334 @@ +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.analysis.MockAnalyzer; +import org.apache.lucene.analysis.MockTokenizer; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.TextField; +import org.apache.lucene.index.*; +import org.apache.lucene.search.*; +import org.apache.lucene.store.Directory; +import org.apache.lucene.util.BytesRef; +import org.apache.lucene.util.FixedBitSet; +import org.apache.lucene.util.LuceneTestCase; +import org.apache.lucene.util._TestUtil; +import org.junit.Test; + +import java.io.IOException; +import java.util.*; + +public class TestJoinUtil extends LuceneTestCase { + + public void testSimple() throws Exception { + final String idField = "id"; + final String toField = "productId"; + + 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", "random text", TextField.TYPE_STORED)); + doc.add(new Field("name", "name1", TextField.TYPE_STORED)); + doc.add(new Field(idField, "1", TextField.TYPE_STORED)); + w.addDocument(doc); + + // 1 + doc = new Document(); + doc.add(new Field("price", "10.0", 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("price", "20.0", 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", "more random text", TextField.TYPE_STORED)); + doc.add(new Field("name", "name2", 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("price", "10.0", 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("price", "20.0", 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 product + MultiTermQuery joinQuery = + JoinUtil.createJoinQuery(idField, toField, new TermQuery(new Term("name", "name2")), indexSearcher); + + 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, toField, new TermQuery(new Term("name", "name1")), indexSearcher); + 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, idField, new TermQuery(new Term("id", "5")), indexSearcher); + result = indexSearcher.search(joinQuery, 10); + assertEquals(1, result.totalHits); + assertEquals(3, result.scoreDocs[0].doc); + + indexSearcher.getIndexReader().close(); + dir.close(); + } + + @Test + public void testRandomJoin() throws Exception { + int maxIndexIter = _TestUtil.nextInt(random, 6, 12); + for (int indexIter = 1; indexIter <= maxIndexIter; indexIter++) { + if (VERBOSE) { + System.out.println("indexIter=" + indexIter); + } + Directory dir = newDirectory(); + RandomIndexWriter w = new RandomIndexWriter( + random, + dir, + newIndexWriterConfig(TEST_VERSION_CURRENT, new MockAnalyzer(random, MockTokenizer.KEYWORD, false)).setMergePolicy(newLogMergePolicy()) + ); + int numberOfDocumentsToIndex = _TestUtil.nextInt(random, 87, 764); + IndexIterationContext context = createContext(numberOfDocumentsToIndex, w); + + IndexReader topLevelReader = w.getReader(); + w.close(); + int maxSearchIter = _TestUtil.nextInt(random, 13, 26); + for (int searchIter = 1; searchIter <= maxSearchIter; searchIter++) { + if (VERBOSE) { + System.out.println("searchIter=" + searchIter); + } + IndexSearcher indexSearcher = newSearcher(topLevelReader); + + int r = random.nextInt(context.randomUniqueValues.length); + boolean from = context.randomFrom[r]; + String randomValue = context.randomUniqueValues[r]; + FixedBitSet expectedResult = createExpectedResult(randomValue, from, indexSearcher.getIndexReader(), context); + + Query actualQuery = new TermQuery(new Term("value", randomValue)); + if (VERBOSE) { + System.out.println("actualQuery=" + actualQuery); + } + MultiTermQuery joinQuery; + if (from) { + joinQuery = JoinUtil.createJoinQuery("from", "to", actualQuery, indexSearcher); + } else { + joinQuery = JoinUtil.createJoinQuery("to", "from", actualQuery, indexSearcher); + } + if (VERBOSE) { + System.out.println("joinQuery=" + joinQuery); + } + + // Need to know all documents that have matches. TopDocs doesn't give me that and then I'd be also testing TopDocsCollector... + final FixedBitSet actualResult = new FixedBitSet(indexSearcher.getIndexReader().maxDoc()); + indexSearcher.search(joinQuery, new Collector() { + int docBase; + + public void setScorer(Scorer scorer) throws IOException { + } + + public void collect(int doc) throws IOException { + actualResult.set(doc + docBase); + } + + public void setNextReader(IndexReader.AtomicReaderContext context) throws IOException { + docBase = context.docBase; + } + + public boolean acceptsDocsOutOfOrder() { + return true; + } + }); + + if (VERBOSE) { + System.out.println("expected cardinality:" + expectedResult.cardinality()); + DocIdSetIterator iterator = expectedResult.iterator(); + for (int doc = iterator.nextDoc(); doc != DocIdSetIterator.NO_MORE_DOCS; doc = iterator.nextDoc()) { + System.out.println(String.format("Expected doc[%d] with id value %s", doc, indexSearcher.doc(doc).get("id"))); + } + System.out.println("actual cardinality:" + actualResult.cardinality()); + iterator = actualResult.iterator(); + for (int doc = iterator.nextDoc(); doc != DocIdSetIterator.NO_MORE_DOCS; doc = iterator.nextDoc()) { + System.out.println(String.format("Actual doc[%d] with id value %s", doc, indexSearcher.doc(doc).get("id"))); + } + } + + assertEquals(expectedResult, actualResult); + } + topLevelReader.close(); + dir.close(); + } + } + + private IndexIterationContext createContext(int nDocs, RandomIndexWriter writer) throws IOException { + return createContext(nDocs, writer, writer); + } + + private IndexIterationContext createContext(int nDocs, RandomIndexWriter fromWriter, RandomIndexWriter toWriter) throws IOException { + IndexIterationContext context = new IndexIterationContext(); + int numRandomValues = nDocs / 2; + context.randomUniqueValues = new String[numRandomValues]; + Set trackSet= new HashSet(); + context.randomFrom = new boolean[numRandomValues]; + for (int i = 0; i < numRandomValues; i++) { + String uniqueRandomValue; + do { + uniqueRandomValue = _TestUtil.randomRealisticUnicodeString(random); + } while ("".equals(uniqueRandomValue) || trackSet.contains(uniqueRandomValue)); + // Generate unique values and empty strings aren't allowed. + trackSet.add(uniqueRandomValue); + context.randomFrom[i] = random.nextBoolean(); + context.randomUniqueValues[i] = uniqueRandomValue; + } + + for (int i = 0; i < nDocs; i++) { + String id = Integer.toString(i); + int randomI = random.nextInt(context.randomUniqueValues.length); + String value = context.randomUniqueValues[randomI]; + String linkValue = context.randomUniqueValues[random.nextInt(context.randomUniqueValues.length)]; + Document document = new Document(); + document.add(newField(random, "id", id, TextField.TYPE_STORED)); + document.add(newField(random, "value", value, TextField.TYPE_STORED)); + + final RandomIndexWriter w; + boolean from = context.randomFrom[randomI]; + if (from) { + if (!context.fromDocuments.containsKey(linkValue)) { + context.fromDocuments.put(linkValue, new ArrayList()); + } + if (!context.randomValueFromDocs.containsKey(value)) { + context.randomValueFromDocs.put(value, new ArrayList()); + } + + RandomDoc doc = new RandomDoc(id, linkValue, value); + context.fromDocuments.get(linkValue).add(doc); + context.randomValueFromDocs.get(value).add(doc); + document.add(newField(random, "from", linkValue, TextField.TYPE_STORED)); + w = fromWriter; + } else { + if (!context.toDocuments.containsKey(linkValue)) { + context.toDocuments.put(linkValue, new ArrayList()); + } + if (!context.randomValueToDocs.containsKey(value)) { + context.randomValueToDocs.put(value, new ArrayList()); + } + + RandomDoc doc = new RandomDoc(id, linkValue, value); + context.toDocuments.get(linkValue).add(doc); + context.randomValueToDocs.get(value).add(doc); + document.add(newField(random, "to", linkValue, TextField.TYPE_STORED)); + w = toWriter; + } + + w.addDocument(document); + if (random.nextInt(10) == 4) { + w.commit(); + } + if (VERBOSE) { + System.out.println("Added document["+ i +"]: " + document); + } + } + return context; + } + + private FixedBitSet createExpectedResult(String queryValue, boolean from, IndexReader topLevelReader, IndexIterationContext context) throws IOException { + final Map> randomValueDocs; + final Map> linkValueDocuments; + if (from) { + randomValueDocs = context.randomValueFromDocs; + linkValueDocuments = context.toDocuments; + } else { + randomValueDocs = context.randomValueToDocs; + linkValueDocuments = context.fromDocuments; + } + + FixedBitSet expectedResult = new FixedBitSet(topLevelReader.maxDoc()); + List matchingDocs = randomValueDocs.get(queryValue); + if (matchingDocs == null) { + return new FixedBitSet(topLevelReader.maxDoc()); + } + + for (RandomDoc matchingDoc : matchingDocs) { + List otherMatchingDocs = linkValueDocuments.get(matchingDoc.link); + if (otherMatchingDocs == null) { + continue; + } + + for (RandomDoc otherSideDoc : otherMatchingDocs) { + DocsEnum docsEnum = MultiFields.getTermDocsEnum(topLevelReader, MultiFields.getLiveDocs(topLevelReader), "id", new BytesRef(otherSideDoc.id), false); + assert docsEnum != null; + int doc = docsEnum.nextDoc(); + expectedResult.set(doc); + } + } + return expectedResult; + } + + private static class IndexIterationContext { + + String[] randomUniqueValues; + boolean[] randomFrom; + Map> fromDocuments = new HashMap>(); + Map> toDocuments = new HashMap>(); + Map> randomValueFromDocs = new HashMap>(); + Map> randomValueToDocs = new HashMap>(); + + } + + private static class RandomDoc { + + final String id; + final String link; + final String value; + + private RandomDoc(String id, String link, String value) { + this.id = id; + this.link = link; + this.value = value; + } + } + +} Index: modules/join/src/java/org/apache/lucene/search/join/TermSetQuery.java IDEA additional info: Subsystem: com.intellij.openapi.diff.impl.patch.CharsetEP <+>UTF-8 =================================================================== --- modules/join/src/java/org/apache/lucene/search/join/TermSetQuery.java (revision ) +++ modules/join/src/java/org/apache/lucene/search/join/TermSetQuery.java (revision ) @@ -0,0 +1,104 @@ +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.FilteredTermsEnum; +import org.apache.lucene.index.Terms; +import org.apache.lucene.index.TermsEnum; +import org.apache.lucene.search.MultiTermQuery; +import org.apache.lucene.util.AttributeSource; +import org.apache.lucene.util.BytesRef; + +import java.io.IOException; +import java.util.NavigableSet; + +/** + * A query that has a set of terms from a specific field. This query will match documents have the same terms as in + * the set. + * + * @lucene.experimental + */ +public class TermSetQuery extends MultiTermQuery { + + private final NavigableSet fromTerms; + private TermsEnum reuse = null; + + /** + * @param fromTerms The set of terms. + * @param toField The field that should contain terms from the set in order for documents to match. + */ + public TermSetQuery(NavigableSet fromTerms, String toField) { + super(toField); + this.fromTerms = fromTerms; + } + + protected TermsEnum getTermsEnum(Terms terms, AttributeSource atts) throws IOException { + return fromTerms.isEmpty() ? TermsEnum.EMPTY : new TermSetTermsEnum(terms.iterator(reuse), fromTerms); + } + + public String toString(String string) { + return "TermSetQuery{" + + "toField=" + field + + '}'; + } + + static class TermSetTermsEnum extends FilteredTermsEnum { + + private final NavigableSet fromTerms; + private BytesRef seekTerm; + + TermSetTermsEnum(TermsEnum tenum, NavigableSet fromTerms) throws IOException { + super(tenum); + this.fromTerms = fromTerms; + seekTerm = fromTerms.first(); + } + + @Override + protected BytesRef nextSeekTerm(BytesRef currentTerm) throws IOException { + BytesRef temp = seekTerm; + seekTerm = null; + return temp; + } + + protected AcceptStatus accept(BytesRef term) throws IOException { + if (term.compareTo(fromTerms.last()) > 0) { + return AcceptStatus.END; + } + + return fromTerms.contains(term) ? AcceptStatus.YES : AcceptStatus.NO; + + // Some quick testing showed me that the above is faster. I guess this depends on how scattered the fromTerms are... + /*if (fromTerms.contains(term)) { + seekTerm = fromTerms.higher(term); + if (seekTerm != null) { + return AcceptStatus.YES_AND_SEEK; + } else { + return AcceptStatus.YES; + } + } else { + seekTerm = fromTerms.higher(term); + if (seekTerm != null) { + return AcceptStatus.NO_AND_SEEK; + } else { + return AcceptStatus.NO; + } + }*/ + } + } + +} Index: modules/join/src/java/org/apache/lucene/search/join/JoinUtil.java IDEA additional info: Subsystem: com.intellij.openapi.diff.impl.patch.CharsetEP <+>UTF-8 =================================================================== --- modules/join/src/java/org/apache/lucene/search/join/JoinUtil.java (revision ) +++ modules/join/src/java/org/apache/lucene/search/join/JoinUtil.java (revision ) @@ -0,0 +1,56 @@ +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.IndexReader; +import org.apache.lucene.search.*; +import org.apache.lucene.util.BytesRef; + +import java.io.IOException; +import java.util.NavigableSet; +import java.util.TreeSet; + +/** + * Utility for query time joining using {@link TermSetQuery}. + * + * @lucene.experimental + */ +public class JoinUtil { + + public static MultiTermQuery createJoinQuery(final String fromField, String toField, Query query, IndexSearcher fromSearcher) throws IOException { + final NavigableSet fromTerms = new TreeSet(); + fromSearcher.search(query, new Collector() { + + private FieldCache.DocTerms fromDocTerms; + + public void collect(int doc) throws IOException { + BytesRef term = fromDocTerms.getTerm(doc, new BytesRef()); + fromTerms.add(term); + } + + public void setNextReader(IndexReader.AtomicReaderContext context) throws IOException { + fromDocTerms = FieldCache.DEFAULT.getTerms(context.reader, fromField); + } + + public void setScorer(Scorer scorer) throws IOException {} + public boolean acceptsDocsOutOfOrder() {return true;} + }); + return new TermSetQuery(fromTerms, toField); + } + +} \ No newline at end of file