Index: suggest/src/test/org/apache/lucene/search/suggest/analyzing/BlendedInfixSuggesterTest.java =================================================================== --- suggest/src/test/org/apache/lucene/search/suggest/analyzing/BlendedInfixSuggesterTest.java (revision 0) +++ suggest/src/test/org/apache/lucene/search/suggest/analyzing/BlendedInfixSuggesterTest.java (revision 0) @@ -0,0 +1,231 @@ +package org.apache.lucene.search.suggest.analyzing; + +/* + * 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.Analyzer; +import org.apache.lucene.analysis.standard.StandardAnalyzer; +import org.apache.lucene.analysis.util.CharArraySet; +import org.apache.lucene.search.suggest.Input; +import org.apache.lucene.search.suggest.InputArrayIterator; +import org.apache.lucene.search.suggest.Lookup; +import org.apache.lucene.store.Directory; +import org.apache.lucene.util.BytesRef; +import org.apache.lucene.util.LuceneTestCase; +import org.apache.lucene.util._TestUtil; + +import java.io.File; +import java.io.IOException; +import java.util.List; + +public class BlendedInfixSuggesterTest extends LuceneTestCase { + + /** + * Test the weight transformation depending on the position + * of the matching term. + */ + public void testBlendedSort() throws IOException { + + BytesRef payload = new BytesRef("star"); + + Input keys[] = new Input[]{ + new Input("star wars: episode v - the empire strikes back", 8, payload) + }; + + File tempDir = _TestUtil.getTempDir("BlendedInfixSuggesterTest"); + + Analyzer a = new StandardAnalyzer(TEST_VERSION_CURRENT, CharArraySet.EMPTY_SET); + BlendedInfixSuggester suggester = new BlendedInfixSuggester(TEST_VERSION_CURRENT, tempDir, a, a, + AnalyzingInfixSuggester.DEFAULT_MIN_PREFIX_CHARS, + BlendedInfixSuggester.BlenderType.POSITION_LINEAR, + BlendedInfixSuggester.DEFAULT_NUM_FACTOR) { + @Override + protected Directory getDirectory(File path) { + return newFSDirectory(path); + } + }; + suggester.build(new InputArrayIterator(keys)); + + // we query for star wars and check that the weight + // is smaller when we search for tokens that are far from the beginning + + long w0 = getInResults(suggester, "star ", payload, 1); + long w1 = getInResults(suggester, "war", payload, 1); + long w2 = getInResults(suggester, "empire ba", payload, 1); + long w3 = getInResults(suggester, "back", payload, 1); + long w4 = getInResults(suggester, "bacc", payload, 1); + + assertTrue(w0 > w1); + assertTrue(w1 > w2); + assertTrue(w2 > w3); + + assertTrue(w4 < 0); + + suggester.close(); + } + + /** + * Verify the different flavours of the blender types + */ + public void testBlendingType() throws IOException { + + BytesRef pl = new BytesRef("lake"); + long w = 20; + + Input keys[] = new Input[]{ + new Input("top of the lake", w, pl) + }; + + File tempDir = _TestUtil.getTempDir("BlendedInfixSuggesterTest"); + Analyzer a = new StandardAnalyzer(TEST_VERSION_CURRENT, CharArraySet.EMPTY_SET); + + // BlenderType.LINEAR is used by default (remove position*10%) + BlendedInfixSuggester suggester = new BlendedInfixSuggester(TEST_VERSION_CURRENT, tempDir, a) { + @Override + protected Directory getDirectory(File path) { + return newFSDirectory(path); + } + }; + suggester.build(new InputArrayIterator(keys)); + + assertEquals(w, getInResults(suggester, "top", pl, 1)); + assertEquals((int) (w * (1 - 0.10 * 2)), getInResults(suggester, "the", pl, 1)); + assertEquals((int) (w * (1 - 0.10 * 3)), getInResults(suggester, "lake", pl, 1)); + + suggester.close(); + + // BlenderType.RECIPROCAL is using 1/(1+p) * w where w is weight and p the position of the word + suggester = new BlendedInfixSuggester(TEST_VERSION_CURRENT, tempDir, a, a, + AnalyzingInfixSuggester.DEFAULT_MIN_PREFIX_CHARS, BlendedInfixSuggester.BlenderType.POSITION_RECIPROCAL, 1) { + @Override + protected Directory getDirectory(File path) { + return newFSDirectory(path); + } + }; + suggester.build(new InputArrayIterator(keys)); + + assertEquals(w, getInResults(suggester, "top", pl, 1)); + assertEquals((int) (w * 1 / (1 + 2)), getInResults(suggester, "the", pl, 1)); + assertEquals((int) (w * 1 / (1 + 3)), getInResults(suggester, "lake", pl, 1)); + + suggester.close(); + } + + /** + * Assert that the factor is important to get results that might be lower in term of weight but + * would be pushed up after the blending transformation + * + * @throws IOException + */ + public void testRequiresMore() throws IOException { + + BytesRef lake = new BytesRef("lake"); + BytesRef star = new BytesRef("star"); + BytesRef ret = new BytesRef("ret"); + + Input keys[] = new Input[]{ + new Input("top of the lake", 15, lake), + new Input("star wars: episode v - the empire strikes back", 12, star), + new Input("the returned", 10, ret), + }; + + File tempDir = _TestUtil.getTempDir("BlendedInfixSuggesterTest"); + Analyzer a = new StandardAnalyzer(TEST_VERSION_CURRENT, CharArraySet.EMPTY_SET); + + // if factor is small, we don't get the expected element + BlendedInfixSuggester suggester = new BlendedInfixSuggester(TEST_VERSION_CURRENT, tempDir, a, a, + AnalyzingInfixSuggester.DEFAULT_MIN_PREFIX_CHARS, BlendedInfixSuggester.BlenderType.POSITION_RECIPROCAL, 1) { + @Override + protected Directory getDirectory(File path) { + return newFSDirectory(path); + } + }; + suggester.build(new InputArrayIterator(keys)); + + long w0 = getInResults(suggester, "the", ret, 2); + assertTrue(w0 < 1); + + long w1 = getInResults(suggester, "the", ret, 3); + assertTrue(w1 > 1); + + suggester.close(); + + // if we increase the factor we have it + suggester = new BlendedInfixSuggester(TEST_VERSION_CURRENT, tempDir, a, a, + AnalyzingInfixSuggester.DEFAULT_MIN_PREFIX_CHARS, BlendedInfixSuggester.BlenderType.POSITION_RECIPROCAL, 2) { + @Override + protected Directory getDirectory(File path) { + return newFSDirectory(path); + } + }; + suggester.build(new InputArrayIterator(keys)); + + long w2 = getInResults(suggester, "the", ret, 2); + assertTrue(w2 > 1); + + long w3 = getInResults(suggester, "the", star, 2); + assertTrue(w3 < 1); + + suggester.close(); + } + + public void /*testT*/rying() throws IOException { + + BytesRef lake = new BytesRef("lake"); + BytesRef star = new BytesRef("star"); + BytesRef ret = new BytesRef("ret"); + + Input keys[] = new Input[]{ + new Input("top of the lake", 15, lake), + new Input("star wars: episode v - the empire strikes back", 12, star), + new Input("the returned", 10, ret), + }; + + File tempDir = _TestUtil.getTempDir("BlendedInfixSuggesterTest"); + Analyzer a = new StandardAnalyzer(TEST_VERSION_CURRENT, CharArraySet.EMPTY_SET); + + // if factor is small, we don't get the expected element + BlendedInfixSuggester suggester = new BlendedInfixSuggester(TEST_VERSION_CURRENT, tempDir, a, a, + AnalyzingInfixSuggester.DEFAULT_MIN_PREFIX_CHARS, BlendedInfixSuggester.BlenderType.POSITION_RECIPROCAL, + BlendedInfixSuggester.DEFAULT_NUM_FACTOR) { + @Override + protected Directory getDirectory(File path) { + return newFSDirectory(path); + } + }; + suggester.build(new InputArrayIterator(keys)); + + + List responses = suggester.lookup("the", 4, true, false); + + for (Lookup.LookupResult response : responses) + System.out.println(response); + + suggester.close(); + } + + private static long getInResults(BlendedInfixSuggester suggester, String prefix, BytesRef payload, int num) { + + List responses = suggester.lookup(prefix, num, true, false); + + for (Lookup.LookupResult response : responses) + if (response.payload.equals(payload)) + return response.value; + + return -1; + } +} Index: suggest/src/java/org/apache/lucene/search/suggest/analyzing/AnalyzingInfixSuggester.java =================================================================== --- suggest/src/java/org/apache/lucene/search/suggest/analyzing/AnalyzingInfixSuggester.java (revision 1552014) +++ suggest/src/java/org/apache/lucene/search/suggest/analyzing/AnalyzingInfixSuggester.java (working copy) @@ -112,9 +112,9 @@ protected IndexSearcher searcher; /** null if payloads were not indexed: */ - private BinaryDocValues payloadsDV; - private BinaryDocValues textDV; - private NumericDocValues weightsDV; + protected BinaryDocValues payloadsDV; + protected BinaryDocValues textDV; + protected NumericDocValues weightsDV; /** Default minimum number of leading characters before * PrefixQuery is used (4). */ @@ -214,9 +214,7 @@ getIndexWriterConfig(matchVersion, gramAnalyzer)); BytesRef text; Document doc = new Document(); - FieldType ft = new FieldType(TextField.TYPE_NOT_STORED); - ft.setIndexOptions(IndexOptions.DOCS_ONLY); - ft.setOmitNorms(true); + FieldType ft = getTextFieldType(); Field textField = new Field(TEXT_FIELD_NAME, "", ft); doc.add(textField); @@ -314,6 +312,18 @@ } } + /** + * Subclass can override this method to change the field type of the text field + * e.g. to change the index options + */ + protected FieldType getTextFieldType(){ + FieldType ft = new FieldType(TextField.TYPE_NOT_STORED); + ft.setIndexOptions(IndexOptions.DOCS_ONLY); + ft.setOmitNorms(true); + + return ft; + } + @Override public List lookup(CharSequence key, boolean onlyMorePopular, int num) { return lookup(key, num, true, true); @@ -413,40 +423,58 @@ // Slower way if postings are not pre-sorted by weight: // hits = searcher.search(query, null, num, new Sort(new SortField("weight", SortField.Type.LONG, true))); - List results = new ArrayList(); - BytesRef scratch = new BytesRef(); - for (int i=0;i results = createResults(hits, num, key, doHighlight, matchedTokens, prefixToken); - BytesRef payload; - if (payloadsDV != null) { - payload = new BytesRef(); - payloadsDV.get(sd.doc, payload); - } else { - payload = null; - } - - LookupResult result; - - if (doHighlight) { - Object highlightKey = highlight(text, matchedTokens, prefixToken); - result = new LookupResult(highlightKey.toString(), highlightKey, score, payload); - } else { - result = new LookupResult(text, score, payload); - } - results.add(result); - } //System.out.println((System.currentTimeMillis() - t0) + " msec for infix suggest"); //System.out.println(results); + return results; + } catch (IOException ioe) { throw new RuntimeException(ioe); } } + /** + * Create the results based on the search hits. + * Can be overridden by subclass to add particular behavior (e.g. weight transformation) + * @throws IOException + */ + protected List createResults(TopDocs hits, int num, CharSequence charSequence, + boolean doHighlight, Set matchedTokens, String prefixToken) + throws IOException { + + List results = new ArrayList(); + BytesRef scratch = new BytesRef(); + for (int i=0;i lookup(CharSequence key, boolean onlyMorePopular, int num) { + // here we multiply the number of searched element by the defined factor + return super.lookup(key, onlyMorePopular, num * numFactor); + } + + @Override + public List lookup(CharSequence key, int num, boolean allTermsRequired, boolean doHighlight) { + // here we multiply the number of searched element by the defined factor + return super.lookup(key, num * numFactor, allTermsRequired, doHighlight); + } + + @Override + protected FieldType getTextFieldType() { + FieldType ft = new FieldType(TextField.TYPE_NOT_STORED); + ft.setIndexOptions(FieldInfo.IndexOptions.DOCS_AND_FREQS_AND_POSITIONS); + ft.setStoreTermVectors(true); + ft.setStoreTermVectorPositions(true); + ft.setOmitNorms(true); + + return ft; + } + + @Override + protected List createResults(TopDocs hits, int num, CharSequence key, + boolean doHighlight, Set matchedTokens, String prefixToken) + throws IOException { + + TreeSet results = new TreeSet<>(LOOKUP_COMP); + + // we reduce the num to the one initially requested + int actualNum = num / numFactor; + + BytesRef scratch = new BytesRef(); + for (int i = 0; i < hits.scoreDocs.length; i++) { + + ScoreDoc sd = hits.scoreDocs[i]; + textDV.get(sd.doc, scratch); + String text = scratch.utf8ToString(); + long weight = weightsDV.get(sd.doc); + + BytesRef payload; + if (payloadsDV != null) { + payload = new BytesRef(); + payloadsDV.get(sd.doc, payload); + } else { + payload = null; + } + + double coefficient; + if (text.startsWith(key.toString())) + // if hit starts with the key, we don't change the score + coefficient = 1; + else + coefficient = createCoefficient(sd.doc, matchedTokens, prefixToken); + + long score = (long) (weight * coefficient); + + LookupResult result; + if (doHighlight) { + Object highlightKey = highlight(text, matchedTokens, prefixToken); + result = new LookupResult(highlightKey.toString(), highlightKey, score, payload); + } else { + result = new LookupResult(text, score, payload); + } + + boundedTreeAdd(results, result, actualNum); + } + + return new ArrayList<>(results.descendingSet()); + } + + /** + * Add an element to the tree respecting a size limit + * + * @param results the tree to add in + * @param result the result we try to add + * @param num size limit + */ + private static void boundedTreeAdd(TreeSet results, Lookup.LookupResult result, int num) { + + if (results.size() >= num) + if (results.first().value < result.value) + results.pollFirst(); + else + return; + + results.add(result); + } + + /** + * Create the coefficient to transform the weight. + * + * @param doc id of the document + * @param matchedTokens tokens found in the query + * @param prefixToken unfinished token in the query + * @return the coefficient + * @throws IOException + */ + private double createCoefficient(int doc, Set matchedTokens, String prefixToken) throws IOException { + + Terms tv = searcher.getIndexReader().getTermVector(doc, TEXT_FIELD_NAME); + TermsEnum it = tv.iterator(TermsEnum.EMPTY); + + Integer position = Integer.MAX_VALUE; + BytesRef term; + // find the closest token position + while ((term = it.next()) != null) { + + String docTerm = term.utf8ToString(); + + if (matchedTokens.contains(docTerm) || docTerm.startsWith(prefixToken)) { + + DocsAndPositionsEnum docPosEnum = it.docsAndPositions(null, null, DocsAndPositionsEnum.FLAG_OFFSETS); + docPosEnum.nextDoc(); + + // use the first occurrence of the term + int p = docPosEnum.nextPosition(); + if (p < position) + position = p; + } + } + + // create corresponding coefficient based on position + double coefficient; + switch (blenderType) { + case POSITION_LINEAR: + coefficient = 1 - LINEAR_COEF * position; + break; + + case POSITION_RECIPROCAL: + coefficient = 1. / (position + 1); + break; + + default: + coefficient = 1; + } + + return coefficient; + } + + private static Comparator LOOKUP_COMP = new LookUpComparator(); + + private static class LookUpComparator implements Comparator { + + @Override + public int compare(Lookup.LookupResult o1, Lookup.LookupResult o2) { + // order on weight + if (o1.value > o2.value) + return 1; + if (o1.value < o2.value) + return 0; + + // otherwise on alphabetic order + return CHARSEQUENCE_COMPARATOR.compare(o1.key, o2.key); + } + } + +} +