refactor(tdigest): convert the flags for reverse in TDigest::Rank to a compile-time argument

src/commands/cmd_tdigest.cc

@@ -176,7 +176,7 @@ class CommandTDigestAdd : public Commander {
   std::vector<double> values_;
 };
 
-template <bool reverse>
+template <bool Reverse>
 class TDigestRankCommand : public Commander {
  public:
   Status Parse(const std::vector<std::string> &args) override {
@@ -202,7 +202,7 @@ class TDigestRankCommand : public Commander {
     TDigest tdigest(srv->storage, conn->GetNamespace());
     std::vector<int> result;
     result.reserve(origin_inputs_.size());
-    if (const auto s = tdigest.Rank(ctx, key_name_, unique_inputs_, reverse, result); !s.ok()) {
+    if (const auto s = tdigest.Rank<Reverse>(ctx, key_name_, unique_inputs_, result); !s.ok()) {
       if (s.IsNotFound()) {
         return {Status::RedisExecErr, errKeyNotFound};
       }

src/types/redis_tdigest.cc

@@ -47,125 +47,6 @@
 
 namespace redis {
 
-// TODO: It should be replaced by a iteration of the rocksdb iterator
-class DummyCentroids {
- public:
-  class BaseIterator {
-   public:
-    virtual ~BaseIterator() = default;
-    virtual bool Next() = 0;
-    virtual bool Prev() = 0;
-    virtual bool Valid() const = 0;
-    virtual std::unique_ptr<BaseIterator> Clone() const = 0;
-    virtual StatusOr<Centroid> GetCentroid() const = 0;
-  };
-
-  DummyCentroids(const TDigestMetadata& meta_data, const std::vector<Centroid>& centroids)
-      : meta_data_(meta_data), centroids_(centroids) {}
-  class Iterator : public BaseIterator {
-   public:
-    Iterator(std::vector<Centroid>::const_iterator&& iter, const std::vector<Centroid>& centroids)
-        : iter_(iter), centroids_(centroids) {}
-    std::unique_ptr<BaseIterator> Clone() const override {
-      if (iter_ != centroids_.cend()) {
-        return std::make_unique<Iterator>(std::next(centroids_.cbegin(), std::distance(centroids_.cbegin(), iter_)),
-                                          centroids_);
-      }
-      return std::make_unique<Iterator>(centroids_.cend(), centroids_);
-    }
-    bool Next() override {
-      if (Valid()) {
-        std::advance(iter_, 1);
-      }
-      return iter_ != centroids_.cend();
-    }
-
-    // The Prev function can only be called for item is not cend,
-    // because we must guarantee the iterator to be inside the valid range before iteration.
-    bool Prev() override {
-      if (Valid() && iter_ != centroids_.cbegin()) {
-        std::advance(iter_, -1);
-      }
-      return Valid();
-    }
-    bool Valid() const override { return iter_ != centroids_.cend(); }
-    StatusOr<Centroid> GetCentroid() const override {
-      if (iter_ == centroids_.cend()) {
-        return {::Status::NotOK, "invalid iterator during decoding tdigest centroid"};
-      }
-      return *iter_;
-    }
-
-   private:
-    std::vector<Centroid>::const_iterator iter_;
-    const std::vector<Centroid>& centroids_;
-  };
-
-  class ReverseIterator final : public BaseIterator {
-   public:
-    ReverseIterator(std::vector<Centroid>::const_reverse_iterator&& iter, const std::vector<Centroid>& centroids)
-        : iter_(iter), centroids_(centroids) {}
-    std::unique_ptr<BaseIterator> Clone() const override {
-      if (iter_ != centroids_.crend()) {
-        return std::make_unique<ReverseIterator>(
-            std::next(centroids_.crbegin(), std::distance(centroids_.crbegin(), iter_)), centroids_);
-      }
-      return std::make_unique<ReverseIterator>(centroids_.crend(), centroids_);
-    }
-    bool Next() override {
-      if (Valid()) {
-        std::advance(iter_, 1);
-      }
-      return iter_ != centroids_.crend();
-    }
-
-    bool Prev() override {
-      if (Valid() && iter_ != centroids_.crbegin()) {
-        std::advance(iter_, -1);
-      }
-      return Valid();
-    }
-    bool Valid() const override { return iter_ != centroids_.crend(); }
-    StatusOr<Centroid> GetCentroid() const override {
-      if (iter_ == centroids_.crend()) {
-        return {::Status::NotOK, "invalid iterator during decoding tdigest centroid"};
-      }
-      return *iter_;
-    }
-
-   private:
-    std::vector<Centroid>::const_reverse_iterator iter_;
-    const std::vector<Centroid>& centroids_;
-  };
-
-  std::unique_ptr<BaseIterator> Begin(const bool reverse = false) const {
-    if (reverse) {
-      return std::make_unique<ReverseIterator>(centroids_.crbegin(), centroids_);
-    }
-    return std::make_unique<Iterator>(centroids_.cbegin(), centroids_);
-  }
-  std::unique_ptr<BaseIterator> End(const bool reverse = false) const {
-    if (centroids_.empty()) {
-      if (reverse) {
-        return std::make_unique<ReverseIterator>(centroids_.crend(), centroids_);
-      }
-      return std::make_unique<Iterator>(centroids_.cend(), centroids_);
-    }
-    if (reverse) {
-      return std::make_unique<ReverseIterator>(std::prev(centroids_.crend()), centroids_);
-    }
-    return std::make_unique<Iterator>(std::prev(centroids_.cend()), centroids_);
-  }
-  double TotalWeight() const { return static_cast<double>(meta_data_.total_weight); }
-  double Min() const { return meta_data_.minimum; }
-  double Max() const { return meta_data_.maximum; }
-  uint64_t Size() const { return meta_data_.merged_nodes; }
-
- private:
-  const TDigestMetadata& meta_data_;
-  const std::vector<Centroid>& centroids_;
-};
-
 uint32_t constexpr kMaxElements = 1 * 1024;  // 1k doubles
 
 rocksdb::Status TDigest::Create(engine::Context& ctx, const Slice& digest_name, const TDigestCreateOptions& options,
@@ -273,40 +154,6 @@ rocksdb::Status TDigest::mergeNodes(engine::Context& ctx, const std::string& ns_
   return rocksdb::Status::OK();
 }
 
-rocksdb::Status TDigest::Rank(engine::Context& ctx, const Slice& digest_name, const std::vector<double>& inputs,
-                              bool reverse, std::vector<int>& result) {
-  auto ns_key = AppendNamespacePrefix(digest_name);
-  TDigestMetadata metadata;
-  {
-    LockGuard guard(storage_->GetLockManager(), ns_key);
-
-    if (auto status = getMetaDataByNsKey(ctx, ns_key, &metadata); !status.ok()) {
-      return status;
-    }
-
-    if (metadata.total_observations == 0) {
-      result.resize(inputs.size(), -2);
-      return rocksdb::Status::OK();
-    }
-
-    if (auto status = mergeNodes(ctx, ns_key, &metadata); !status.ok()) {
-      return status;
-    }
-  }
-
-  std::vector<Centroid> centroids;
-  if (auto status = dumpCentroids(ctx, ns_key, metadata, &centroids); !status.ok()) {
-    return status;
-  }
-
-  auto dump_centroids = DummyCentroids(metadata, centroids);
-  auto status = TDigestRank(dump_centroids, inputs, reverse, result);
-  if (!status) {
-    return rocksdb::Status::InvalidArgument(status.Msg());
-  }
-  return rocksdb::Status::OK();
-}
-
 rocksdb::Status TDigest::Quantile(engine::Context& ctx, const Slice& digest_name, const std::vector<double>& qs,
                                   TDigestQuantitleResult* result) {
   auto ns_key = AppendNamespacePrefix(digest_name);
@@ -332,7 +179,7 @@ rocksdb::Status TDigest::Quantile(engine::Context& ctx, const Slice& digest_name
     return status;
   }
 
-  auto dump_centroids = DummyCentroids(metadata, centroids);
+  auto dump_centroids = DummyCentroids<false>(metadata, centroids);
 
   auto quantile_results = std::vector<double>();
   quantile_results.reserve(qs.size());

src/types/redis_tdigest.h

@@ -33,6 +33,101 @@
 #include "tdigest.h"
 
 namespace redis {
+
+// TODO: It should be replaced by a iteration of the rocksdb iterator
+template <bool Reverse>
+class DummyCentroids {
+ public:
+  DummyCentroids(const TDigestMetadata& meta_data, const std::vector<Centroid>& centroids)
+      : meta_data_(meta_data), centroids_(centroids) {}
+  class Iterator {
+   public:
+    using IterType = std::conditional_t<Reverse, std::vector<Centroid>::const_reverse_iterator,
+                                        std::vector<Centroid>::const_iterator>;
+    Iterator(IterType iter, const std::vector<Centroid>& centroids) : iter_(iter), centroids_(centroids) {}
+    std::unique_ptr<Iterator> Clone() const {
+      if (iter_ != getCendIter(centroids_)) {
+        return std::make_unique<Iterator>(
+            std::next(getCbeginIter(centroids_), std::distance(getCbeginIter(centroids_), iter_)), centroids_);
+      }
+      return std::make_unique<Iterator>(getCendIter(centroids_), centroids_);
+    }
+    bool Next() {
+      if (Valid()) {
+        std::advance(iter_, 1);
+      }
+      return iter_ != getCendIter(centroids_);
+    }
+
+    // The Prev function can only be called for item is not cend,
+    // because we must guarantee the iterator to be inside the valid range before iteration.
+    bool Prev() {
+      if (Valid() && iter_ != getCendIter(centroids_)) {
+        std::advance(iter_, -1);
+      }
+      return Valid();
+    }
+    bool Valid() const { return iter_ != getCendIter(centroids_); }
+    StatusOr<Centroid> GetCentroid() const {
+      if (iter_ == getCendIter(centroids_)) {
+        return {::Status::NotOK, "invalid iterator during decoding tdigest centroid"};
+      }
+      return *iter_;
+    }
+
+   private:
+    IterType iter_;
+    const std::vector<Centroid>& centroids_;
+    template <typename Container>
+    decltype(auto) getCbeginIter(const Container& centroids) const {
+      if constexpr (Reverse) {
+        return centroids.crbegin();
+      } else {
+        return centroids.cbegin();
+      }
+    }
+
+    template <typename Container>
+    decltype(auto) getCendIter(const Container& centroids) const {
+      if constexpr (Reverse) {
+        return centroids.crend();
+      } else {
+        return centroids.cend();
+      }
+    }
+  };
+
+  std::unique_ptr<Iterator> Begin() const {
+    if constexpr (Reverse) {
+      return std::make_unique<Iterator>(centroids_.crbegin(), centroids_);
+    } else {
+      return std::make_unique<Iterator>(centroids_.cbegin(), centroids_);
+    }
+  }
+  std::unique_ptr<Iterator> End() const {
+    if (centroids_.empty()) {
+      if constexpr (Reverse) {
+        return std::make_unique<Iterator>(centroids_.crend(), centroids_);
+      } else {
+        return std::make_unique<Iterator>(centroids_.cend(), centroids_);
+      }
+    }
+    if constexpr (Reverse) {
+      return std::make_unique<Iterator>(std::prev(centroids_.crend()), centroids_);
+    } else {
+      return std::make_unique<Iterator>(std::prev(centroids_.cend()), centroids_);
+    }
+  }
+  double TotalWeight() const { return static_cast<double>(meta_data_.total_weight); }
+  double Min() const { return meta_data_.minimum; }
+  double Max() const { return meta_data_.maximum; }
+  uint64_t Size() const { return meta_data_.merged_nodes; }
+
+ private:
+  const TDigestMetadata& meta_data_;
+  const std::vector<Centroid>& centroids_;
+};
+
 inline constexpr uint32_t kTDigestMaxCompression = 1000;  // limit the compression to 1k
 
 struct CentroidWithKey {
@@ -77,7 +172,8 @@ class TDigest : public SubKeyScanner {
 
   rocksdb::Status Merge(engine::Context& ctx, const Slice& dest_digest, const std::vector<std::string>& source_digests,
                         const TDigestMergeOptions& options);
-  rocksdb::Status Rank(engine::Context& ctx, const Slice& digest_name, const std::vector<double>& inputs, bool reverse,
+  template <bool Reverse>
+  rocksdb::Status Rank(engine::Context& ctx, const Slice& digest_name, const std::vector<double>& inputs,
                        std::vector<int>& result);
   rocksdb::Status GetMetaData(engine::Context& context, const Slice& digest_name, TDigestMetadata* metadata);
 
@@ -132,4 +228,38 @@ class TDigest : public SubKeyScanner {
                                              Centroid* centroid) const;
 };
 
+template <bool Reverse>
+rocksdb::Status TDigest::Rank(engine::Context& ctx, const Slice& digest_name, const std::vector<double>& inputs,
+                              std::vector<int>& result) {
+  auto ns_key = AppendNamespacePrefix(digest_name);
+  TDigestMetadata metadata;
+  {
+    LockGuard guard(storage_->GetLockManager(), ns_key);
+
+    if (auto status = getMetaDataByNsKey(ctx, ns_key, &metadata); !status.ok()) {
+      return status;
+    }
+
+    if (metadata.total_observations == 0) {
+      result.resize(inputs.size(), -2);
+      return rocksdb::Status::OK();
+    }
+
+    if (auto status = mergeNodes(ctx, ns_key, &metadata); !status.ok()) {
+      return status;
+    }
+  }
+
+  std::vector<Centroid> centroids;
+  if (auto status = dumpCentroids(ctx, ns_key, metadata, &centroids); !status.ok()) {
+    return status;
+  }
+
+  auto dump_centroids = DummyCentroids<Reverse>(metadata, centroids);
+  if (auto status = TDigestRank<Reverse>(dump_centroids, inputs, result); !status) {
+    return rocksdb::Status::InvalidArgument(status.Msg());
+  }
+  return rocksdb::Status::OK();
+}
+
 }  // namespace redis

src/types/tdigest.h

@@ -171,8 +171,8 @@ struct DoubleComparator {
   bool operator()(const double& a, const double& b) const { return DoubleCompare(a, b) == -1; }
 };
 
-template <typename TD, bool Reverse>
-inline Status TDigestRankImpl(TD&& td, const std::vector<double>& inputs, std::vector<int>& result) {
+template <bool Reverse, typename TD>
+inline Status TDigestRank(TD&& td, const std::vector<double>& inputs, std::vector<int>& result) {
   std::map<double, size_t, DoubleComparator> value_to_index;
   for (size_t i = 0; i < inputs.size(); ++i) {
     value_to_index[inputs[i]] = i;
@@ -211,7 +211,7 @@ inline Status TDigestRankImpl(TD&& td, const std::vector<double>& inputs, std::v
     }
   }
 
-  auto iter = td.Begin(Reverse);
+  auto iter = td.Begin();
   double cumulative_weight = 0;
   while (iter->Valid() && !is_end()) {
     auto centroid = GET_OR_RET(iter->GetCentroid());
@@ -267,12 +267,3 @@ inline Status TDigestRankImpl(TD&& td, const std::vector<double>& inputs, std::v
   }
   return Status::OK();
 }
-
-template <typename TD>
-inline Status TDigestRank(TD&& td, const std::vector<double>& inputs, bool reverse, std::vector<int>& result) {
-  if (reverse) {
-    return TDigestRankImpl<TD, true>(std::forward<TD>(td), inputs, result);
-  } else {
-    return TDigestRankImpl<TD, false>(std::forward<TD>(td), inputs, result);
-  }
-}