Prefix completions with custom python filter

keyvi/include/keyvi/dictionary/completion/forward_backward_completion.h

@@ -53,7 +53,7 @@ class ForwardBackwardCompletion final {
   ForwardBackwardCompletion(dictionary_t forward_dictionary, dictionary_t backward_dictionary)
       : forward_completions_(forward_dictionary), backward_completions_(backward_dictionary) {}
 
-  struct result_compare : public std::binary_function<Match, Match, bool> {
+  struct result_compare {
     bool operator()(const Match& m1, const Match& m2) const { return m1.GetScore() < m2.GetScore(); }
   };
 

keyvi/include/keyvi/dictionary/dictionary.h

@@ -38,6 +38,8 @@
 #include "keyvi/dictionary/match_iterator.h"
 #include "keyvi/dictionary/matching/fuzzy_matching.h"
 #include "keyvi/dictionary/matching/near_matching.h"
+#include "keyvi/dictionary/matching/prefix_completion_matching.h"
+#include "keyvi/dictionary/util/bounded_priority_queue.h"
 
 // #define ENABLE_TRACING
 #include "keyvi/dictionary/util/trace.h"
@@ -324,6 +326,40 @@ class Dictionary final {
     return MatchIterator::MakeIteratorPair(func, data->FirstMatch());
   }
 
+  MatchIterator::MatchIteratorPair GetPrefixCompletion(const std::string& query) const {
+    auto data = std::make_shared<matching::PrefixCompletionMatching<>>(
+        matching::PrefixCompletionMatching<>::FromSingleFsa(fsa_, query));
+
+    auto func = [data]() { return data->NextMatch(); };
+    return MatchIterator::MakeIteratorPair(
+        func, data->FirstMatch(),
+        std::bind(&matching::PrefixCompletionMatching<>::SetMinWeight, &(*data), std::placeholders::_1));
+  }
+
+  MatchIterator::MatchIteratorPair GetPrefixCompletion(const std::string& query, size_t top_n) const {
+    auto data = std::make_shared<matching::PrefixCompletionMatching<>>(
+        matching::PrefixCompletionMatching<>::FromSingleFsa(fsa_, query));
+
+    auto best_weights = std::make_shared<util::BoundedPriorityQueue<uint32_t>>(top_n);
+
+    auto func = [data, best_weights = std::move(best_weights)]() {
+      auto m = data->NextMatch();
+      while (!m.IsEmpty()) {
+        if (m.GetWeight() >= best_weights->Back()) {
+          best_weights->Put(m.GetWeight());
+          return m;
+        }
+
+        m = data->NextMatch();
+      }
+      return Match();
+    };
+
+    return MatchIterator::MakeIteratorPair(
+        func, data->FirstMatch(),
+        std::bind(&matching::PrefixCompletionMatching<>::SetMinWeight, &(*data), std::placeholders::_1));
+  }
+
   std::string GetManifest() const { return fsa_->GetManifest(); }
 
  private:

keyvi/include/keyvi/dictionary/fsa/comparable_state_traverser.h

@@ -168,6 +168,15 @@ class ComparableStateTraverser final {
 
   size_t GetOrder() const { return order_; }
 
+  /**
+   * Set the minimum weight states must be greater or equal to.
+   *
+   * Only available for WeightedTransition specialization.
+   *
+   * @param weight minimum transition weight
+   */
+  inline void SetMinWeight(uint32_t weight) {}
+
  private:
   innerTraverserType state_traverser_;
   std::vector<label_t> label_stack_;
@@ -267,6 +276,16 @@ inline bool ComparableStateTraverser<NearStateTraverser>::operator<(const Compar
   return order_ > rhs.order_;
 }
 
+/**
+ * Set the minimum weight states must be greater or equal to.
+ *
+ * @param weight minimum transition weight
+ */
+template <>
+inline void ComparableStateTraverser<WeightedStateTraverser>::SetMinWeight(uint32_t weight) {
+  state_traverser_.SetMinWeight(weight);
+}
+
 } /* namespace fsa */
 } /* namespace dictionary */
 } /* namespace keyvi */

keyvi/include/keyvi/dictionary/fsa/state_traverser.h

@@ -29,6 +29,7 @@
 
 #include "keyvi/dictionary/fsa/automata.h"
 #include "keyvi/dictionary/fsa/traversal/traversal_base.h"
+#include "keyvi/dictionary/fsa/traversal/weighted_traversal.h"
 
 // #define ENABLE_TRACING
 #include "keyvi/dictionary/util/trace.h"
@@ -119,13 +120,14 @@ class StateTraverser final {
 
   void operator++(int) {
     TRACE("statetraverser++");
+
     // ignore cases where we are already at the end
     if (current_state_ == 0) {
       TRACE("at the end");
       return;
     }
 
-    current_state_ = stack_.GetStates().GetNextState();
+    current_state_ = FilterByMinWeight(stack_.GetStates().GetNextState());
     TRACE("next state: %ld depth: %ld", current_state_, stack_.GetDepth());
 
     while (current_state_ == 0) {
@@ -139,7 +141,7 @@ class StateTraverser final {
       TRACE("state is 0, go up");
       --stack_;
       stack_.GetStates()++;
-      current_state_ = stack_.GetStates().GetNextState();
+      current_state_ = FilterByMinWeight(stack_.GetStates().GetNextState());
       TRACE("next state %ld depth %ld", current_state_, stack_.GetDepth());
     }
 
@@ -155,6 +157,15 @@ class StateTraverser final {
 
   operator bool() const { return !at_end_; }
 
+  /**
+   * Set the minimum weight states must be greater or equal to.
+   *
+   * Only available for WeightedTransition specialization.
+   *
+   * @param min_weight minimum transition weight
+   */
+  inline void SetMinWeight(uint32_t min_weight) {}
+
   bool AtEnd() const { return at_end_; }
 
  private:
@@ -165,6 +176,13 @@ class StateTraverser final {
   bool at_end_;
   traversal::TraversalStack<TransitionT> stack_;
 
+  /**
+   * Filter hook for weighted traversal to filter weights lower than the minimum weight (see spezialisation).
+   *
+   * Default: no filter
+   */
+  inline uint64_t FilterByMinWeight(uint64_t state) { return state; }
+
   template <class innerTraverserType>
   friend class ComparableStateTraverser;
   const traversal::TraversalStack<TransitionT> &GetStack() const { return stack_; }
@@ -176,6 +194,32 @@ class StateTraverser final {
   const traversal::TraversalPayload<TransitionT> &GetTraversalPayload() const { return stack_.traversal_stack_payload; }
 };
 
+/**
+ * Filter state that doesn't meet the min weight requirement.
+ *
+ * This happens when SetMinWeight has been calles after the transitions got already read.
+ *
+ * @param state the state we currently look at.
+ *
+ * @return the current state if it has a weight higher than the minimum weight, 0 otherwise.
+ */
+template <>
+inline uint64_t StateTraverser<traversal::WeightedTransition>::FilterByMinWeight(uint64_t state) {
+  TRACE("filter min weight for weighted transition specialization");
+  return state > 0 && stack_.GetStates().GetNextInnerWeight() >= stack_.traversal_stack_payload.min_weight ? state : 0;
+}
+
+/**
+ * Set the minimum weight states must be greater or equal to.
+ *
+ * @param weight minimum transition weight
+ */
+template <>
+inline void StateTraverser<traversal::WeightedTransition>::SetMinWeight(uint32_t min_weight) {
+  TRACE("set min weight for weighted transition specialization %d", min_weight);
+  stack_.traversal_stack_payload.min_weight = min_weight;
+}
+
 } /* namespace fsa */
 } /* namespace dictionary */
 } /* namespace keyvi */

keyvi/include/keyvi/dictionary/fsa/traversal/weighted_traversal.h

@@ -52,6 +52,20 @@ static bool WeightedTransitionCompare(const WeightedTransition& a, const Weighte
   return a.weight > b.weight;
 }
 
+template <>
+struct TraversalPayload<WeightedTransition> {
+  size_t current_depth;
+  uint32_t min_weight = 0;
+};
+
+template <>
+inline void TraversalState<WeightedTransition>::Add(uint64_t s, uint32_t weight, unsigned char l,
+                                                    TraversalPayload<WeightedTransition>* payload) {
+  if (weight >= payload->min_weight) {
+    traversal_state_payload.transitions.push_back(WeightedTransition(s, weight, l));
+  }
+}
+
 template <>
 inline void TraversalState<WeightedTransition>::PostProcess(TraversalPayload<WeightedTransition>* payload) {
   if (traversal_state_payload.transitions.size() > 0) {

keyvi/include/keyvi/dictionary/fsa/traverser_types.h

@@ -26,15 +26,13 @@
 #define KEYVI_DICTIONARY_FSA_TRAVERSER_TYPES_H_
 
 #include "keyvi/dictionary/fsa/state_traverser.h"
-#include "keyvi/dictionary/fsa/traversal/bounded_weighted_traversal.h"
 #include "keyvi/dictionary/fsa/traversal/near_traversal.h"
 #include "keyvi/dictionary/fsa/traversal/weighted_traversal.h"
 
 namespace keyvi {
 namespace dictionary {
 namespace fsa {
 
-using BoundedWeightedStateTraverser2 = StateTraverser<traversal::BoundedWeightedTransition>;
 using NearStateTraverser = StateTraverser<traversal::NearTransition>;
 using WeightedStateTraverser = StateTraverser<traversal::WeightedTransition>;
 

keyvi/include/keyvi/dictionary/fsa/zip_state_traverser.h

@@ -214,6 +214,15 @@ class ZipStateTraverser final {
 
   const std::vector<label_t> &GetStateLabels() const { return traverser_queue_.top()->GetStateLabels(); }
 
+  /**
+   * Set the minimum weight states must be greater or equal to.
+   *
+   * Only available for WeightedTransition specialization.
+   *
+   * @param weight minimum transition weight
+   */
+  inline void SetMinWeight(uint32_t weight) {}
+
  private:
   heap_t traverser_queue_;
   bool final_ = false;
@@ -226,6 +235,8 @@ class ZipStateTraverser final {
   automata_t fsa_;
   size_t equal_states_ = 1;
   bool pruned = false;
+  // this field only used for weighted traversal, ignored otherwise
+  uint32_t min_weight_ = 0;
 
   inline void PreIncrement() {}
 
@@ -500,6 +511,18 @@ inline ZipStateTraverser<WeightedStateTraverser>::ZipStateTraverser(
   FillInValues();
 }
 
+/**
+ * Set the minimum weight states must be greater or equal to.
+ *
+ * @param weight minimum transition weight
+ */
+template <>
+inline void ZipStateTraverser<WeightedStateTraverser>::SetMinWeight(uint32_t weight) {
+  for (auto t : traverser_queue_) {
+    t->SetMinWeight(weight);
+  }
+}
+
 }  // namespace fsa
 }  // namespace dictionary
 }  // namespace keyvi