/*
* 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.
*/
package org.apache.lucene.search.suggest.fst;
import java.io.IOException;
import java.util.*;
import org.apache.lucene.util.*;
import org.apache.lucene.util.fst.FST;
import org.apache.lucene.util.fst.FST.Arc;
/**
* Finite state automata based implementation of "autocomplete" functionality.
*
* @see FSTCompletionBuilder
* @lucene.experimental
*/
// TODO: we could store exact weights as outputs from the FST (int4 encoded
// floats). This would provide exact outputs from this method and to some
// degree allowed post-sorting on a more fine-grained weight.
// TODO: support for Analyzers (infix suggestions, synonyms?)
public class FSTCompletion {
/**
* A single completion for a given key.
*/
public static final class Completion implements Comparable<Completion> {
/** UTF-8 bytes of the suggestion */
public final BytesRef utf8;
/** source bucket (weight) of the suggestion */
public final int bucket;
Completion(BytesRef key, int bucket) {
this.utf8 = BytesRef.deepCopyOf(key);
this.bucket = bucket;
}
@Override
public String toString() {
return utf8.utf8ToString() + "/" + bucket;
}
/** @see BytesRef#compareTo(BytesRef) */
@Override
public int compareTo(Completion o) {
return this.utf8.compareTo(o.utf8);
}
}
/**
* Default number of buckets.
*/
public static final int DEFAULT_BUCKETS = 10;
/**
* An empty result. Keep this an {@link ArrayList} to keep all the returned
* lists of single type (monomorphic calls).
*/
private static final ArrayList<Completion> EMPTY_RESULT = new ArrayList<>();
/**
* Finite state automaton encoding all the lookup terms. See class notes for
* details.
*/
private final FST<Object> automaton;
/**
* An array of arcs leaving the root automaton state and encoding weights of
* all completions in their sub-trees.
*/
private final Arc<Object>[] rootArcs;
/**
* @see #FSTCompletion(FST, boolean, boolean)
*/
private boolean exactFirst;
/**
* @see #FSTCompletion(FST, boolean, boolean)
*/
private boolean higherWeightsFirst;
/**
* Constructs an FSTCompletion, specifying higherWeightsFirst and exactFirst.
* @param automaton
* Automaton with completions. See {@link FSTCompletionBuilder}.
* @param higherWeightsFirst
* Return most popular suggestions first. This is the default
* behavior for this implementation. Setting it to <code>false</code>
* has no effect (use constant term weights to sort alphabetically
* only).
* @param exactFirst
* Find and push an exact match to the first position of the result
* list if found.
*/
@SuppressWarnings({"unchecked","rawtypes"})
public FSTCompletion(FST<Object> automaton, boolean higherWeightsFirst, boolean exactFirst) {
this.automaton = automaton;
if (automaton != null) {
this.rootArcs = cacheRootArcs(automaton);
} else {
this.rootArcs = new Arc[0];
}
this.higherWeightsFirst = higherWeightsFirst;
this.exactFirst = exactFirst;
}
/**
* Defaults to higher weights first and exact first.
* @see #FSTCompletion(FST, boolean, boolean)
*/
public FSTCompletion(FST<Object> automaton) {
this(automaton, true, true);
}
/**
* Cache the root node's output arcs starting with completions with the
* highest weights.
*/
@SuppressWarnings({"unchecked","rawtypes"})
private static Arc<Object>[] cacheRootArcs(FST<Object> automaton) {
try {
List<Arc<Object>> rootArcs = new ArrayList<>();
Arc<Object> arc = automaton.getFirstArc(new Arc<>());
FST.BytesReader fstReader = automaton.getBytesReader();
automaton.readFirstTargetArc(arc, arc, fstReader);
while (true) {
rootArcs.add(new Arc<>().copyFrom(arc));
if (arc.isLast()) break;
automaton.readNextArc(arc, fstReader);
}
Collections.reverse(rootArcs); // we want highest weights first.
return rootArcs.toArray(new Arc[rootArcs.size()]);
} catch (IOException e) {
throw new RuntimeException(e);
}
}
/**
* Returns the first exact match by traversing root arcs, starting from the
* arc <code>rootArcIndex</code>.
*
* @param rootArcIndex
* The first root arc index in {@link #rootArcs} to consider when
* matching.
*
* @param utf8
* The sequence of utf8 bytes to follow.
*
* @return Returns the bucket number of the match or <code>-1</code> if no
* match was found.
*/
private int getExactMatchStartingFromRootArc(
int rootArcIndex, BytesRef utf8) {
// Get the UTF-8 bytes representation of the input key.
try {
final FST.Arc<Object> scratch = new FST.Arc<>();
FST.BytesReader fstReader = automaton.getBytesReader();
for (; rootArcIndex < rootArcs.length; rootArcIndex++) {
final FST.Arc<Object> rootArc = rootArcs[rootArcIndex];
final FST.Arc<Object> arc = scratch.copyFrom(rootArc);
// Descend into the automaton using the key as prefix.
if (descendWithPrefix(arc, utf8)) {
automaton.readFirstTargetArc(arc, arc, fstReader);
if (arc.label == FST.END_LABEL) {
// Normalize prefix-encoded weight.
return rootArc.label;
}
}
}
} catch (IOException e) {
// Should never happen, but anyway.
throw new RuntimeException(e);
}
// No match.
return -1;
}
/**
* Lookup suggestions to <code>key</code>.
*
* @param key
* The prefix to which suggestions should be sought.
* @param num
* At most this number of suggestions will be returned.
* @return Returns the suggestions, sorted by their approximated weight first
* (decreasing) and then alphabetically (UTF-8 codepoint order).
*/
public List<Completion> lookup(CharSequence key, int num) {
if (key.length() == 0 || automaton == null) {
return EMPTY_RESULT;
}
try {
BytesRef keyUtf8 = new BytesRef(key);
if (!higherWeightsFirst && rootArcs.length > 1) {
// We could emit a warning here (?). An optimal strategy for
// alphabetically sorted
// suggestions would be to add them with a constant weight -- this saves
// unnecessary
// traversals and sorting.
return lookupSortedAlphabetically(keyUtf8, num);
} else {
return lookupSortedByWeight(keyUtf8, num, false);
}
} catch (IOException e) {
// Should never happen, but anyway.
throw new RuntimeException(e);
}
}
/**
* Lookup suggestions sorted alphabetically <b>if weights are not
* constant</b>. This is a workaround: in general, use constant weights for
* alphabetically sorted result.
*/
private List<Completion> lookupSortedAlphabetically(BytesRef key, int num)
throws IOException {
// Greedily get num results from each weight branch.
List<Completion> res = lookupSortedByWeight(key, num, true);
// Sort and trim.
Collections.sort(res);
if (res.size() > num) {
res = res.subList(0, num);
}
return res;
}
/**
* Lookup suggestions sorted by weight (descending order).
*
* @param collectAll
* If <code>true</code>, the routine terminates immediately when
* <code>num</code> suggestions have been collected. If
* <code>false</code>, it will collect suggestions from all weight
* arcs (needed for {@link #lookupSortedAlphabetically}.
*/
private ArrayList<Completion> lookupSortedByWeight(BytesRef key,
int num, boolean collectAll) throws IOException {
// Don't overallocate the results buffers. This also serves the purpose of
// allowing the user of this class to request all matches using Integer.MAX_VALUE as
// the number of results.
final ArrayList<Completion> res = new ArrayList<>(Math.min(10, num));
final BytesRef output = BytesRef.deepCopyOf(key);
for (int i = 0; i < rootArcs.length; i++) {
final FST.Arc<Object> rootArc = rootArcs[i];
final FST.Arc<Object> arc = new FST.Arc<>().copyFrom(rootArc);
// Descend into the automaton using the key as prefix.
if (descendWithPrefix(arc, key)) {
// A subgraph starting from the current node has the completions
// of the key prefix. The arc we're at is the last key's byte,
// so we will collect it too.
output.length = key.length - 1;
if (collect(res, num, rootArc.label, output, arc) && !collectAll) {
// We have enough suggestions to return immediately. Keep on looking
// for an
// exact match, if requested.
if (exactFirst) {
if (!checkExistingAndReorder(res, key)) {
int exactMatchBucket = getExactMatchStartingFromRootArc(i, key);
if (exactMatchBucket != -1) {
// Insert as the first result and truncate at num.
while (res.size() >= num) {
res.remove(res.size() - 1);
}
res.add(0, new Completion(key, exactMatchBucket));
}
}
}
break;
}
}
}
return res;
}
/**
* Checks if the list of
* {@link org.apache.lucene.search.suggest.Lookup.LookupResult}s already has a
* <code>key</code>. If so, reorders that
* {@link org.apache.lucene.search.suggest.Lookup.LookupResult} to the first
* position.
*
* @return Returns <code>true<code> if and only if <code>list</code> contained
* <code>key</code>.
*/
private boolean checkExistingAndReorder(ArrayList<Completion> list, BytesRef key) {
// We assume list does not have duplicates (because of how the FST is created).
for (int i = list.size(); --i >= 0;) {
if (key.equals(list.get(i).utf8)) {
// Key found. Unless already at i==0, remove it and push up front so
// that the ordering
// remains identical with the exception of the exact match.
list.add(0, list.remove(i));
return true;
}
}
return false;
}
/**
* Descend along the path starting at <code>arc</code> and going through bytes
* in the argument.
*
* @param arc
* The starting arc. This argument is modified in-place.
* @param utf8
* The term to descend along.
* @return If <code>true</code>, <code>arc</code> will be set to the arc
* matching last byte of <code>term</code>. <code>false</code> is
* returned if no such prefix exists.
*/
private boolean descendWithPrefix(Arc<Object> arc, BytesRef utf8)
throws IOException {
final int max = utf8.offset + utf8.length;
// Cannot save as instance var since multiple threads
// can use FSTCompletion at once...
final FST.BytesReader fstReader = automaton.getBytesReader();
for (int i = utf8.offset; i < max; i++) {
if (automaton.findTargetArc(utf8.bytes[i] & 0xff, arc, arc, fstReader) == null) {
// No matching prefixes, return an empty result.
return false;
}
}
return true;
}
/**
* Recursive collect lookup results from the automaton subgraph starting at
* <code>arc</code>.
*
* @param num
* Maximum number of results needed (early termination).
*/
private boolean collect(List<Completion> res, int num, int bucket,
BytesRef output, Arc<Object> arc) throws IOException {
if (output.length == output.bytes.length) {
output.bytes = ArrayUtil.grow(output.bytes);
}
assert output.offset == 0;
output.bytes[output.length++] = (byte) arc.label;
FST.BytesReader fstReader = automaton.getBytesReader();
automaton.readFirstTargetArc(arc, arc, fstReader);
while (true) {
if (arc.label == FST.END_LABEL) {
res.add(new Completion(output, bucket));
if (res.size() >= num) return true;
} else {
int save = output.length;
if (collect(res, num, bucket, output, new Arc<>().copyFrom(arc))) {
return true;
}
output.length = save;
}
if (arc.isLast()) {
break;
}
automaton.readNextArc(arc, fstReader);
}
return false;
}
/**
* Returns the bucket count (discretization thresholds).
*/
public int getBucketCount() {
return rootArcs.length;
}
/**
* Returns the bucket assigned to a given key (if found) or <code>-1</code> if
* no exact match exists.
*/
public int getBucket(CharSequence key) {
return getExactMatchStartingFromRootArc(0, new BytesRef(key));
}
/**
* Returns the internal automaton.
*/
public FST<Object> getFST() {
return automaton;
}
}