feat(tdigest): add the support of TDIGEST.REVRANK command

src/commands/cmd_tdigest.cc

@@ -176,6 +176,54 @@ class CommandTDigestAdd : public Commander {
   std::vector<double> values_;
 };
 
+class CommandTDigestRevRank : public Commander {
+ public:
+  Status Parse(const std::vector<std::string> &args) override {
+    key_name_ = args[1];
+
+    std::set<std::string> unique_inputs_set(args.begin() + 2, args.end());
+    origin_inputs_.assign(args.begin() + 2, args.end());
+
+    unique_inputs_.reserve(unique_inputs_set.size());
+    size_t i = 0;
+    for (const auto &input : unique_inputs_set) {
+      auto value = ParseFloat(input);
+      if (!value) {
+        return {Status::RedisParseErr, errValueIsNotFloat};
+      }
+      unique_inputs_.push_back(*value);
+      unique_inputs_order_[input] = i;
+      ++i;
+    }
+    return Status::OK();
+  }
+  Status Execute(engine::Context &ctx, Server *srv, Connection *conn, std::string *output) override {
+    TDigest tdigest(srv->storage, conn->GetNamespace());
+    std::vector<int> result;
+    result.reserve(origin_inputs_.size());
+    if (const auto s = tdigest.RevRank(ctx, key_name_, unique_inputs_, result); !s.ok()) {
+      if (s.IsNotFound()) {
+        return {Status::RedisExecErr, errKeyNotFound};
+      }
+      return {Status::RedisExecErr, s.ToString()};
+    }
+
+    std::vector<std::string> rev_ranks;
+    rev_ranks.reserve(origin_inputs_.size());
+    for (const auto &v : origin_inputs_) {
+      rev_ranks.push_back(redis::Integer(result[unique_inputs_order_[v]]));
+    }
+    *output = redis::Array(rev_ranks);
+    return Status::OK();
+  }
+
+ private:
+  std::string key_name_;
+  std::vector<double> unique_inputs_;
+  std::map<std::string, size_t> unique_inputs_order_;
+  std::vector<std::string> origin_inputs_;
+};
+
 class CommandTDigestMinMax : public Commander {
  public:
   explicit CommandTDigestMinMax(bool is_min) : is_min_(is_min) {}
@@ -369,6 +417,7 @@ REDIS_REGISTER_COMMANDS(TDigest, MakeCmdAttr<CommandTDigestCreate>("tdigest.crea
                         MakeCmdAttr<CommandTDigestAdd>("tdigest.add", -3, "write", 1, 1, 1),
                         MakeCmdAttr<CommandTDigestMax>("tdigest.max", 2, "read-only", 1, 1, 1),
                         MakeCmdAttr<CommandTDigestMin>("tdigest.min", 2, "read-only", 1, 1, 1),
+                        MakeCmdAttr<CommandTDigestRevRank>("tdigest.revrank", -3, "read-only", 1, 1, 1),
                         MakeCmdAttr<CommandTDigestQuantile>("tdigest.quantile", -3, "read-only", 1, 1, 1),
                         MakeCmdAttr<CommandTDigestReset>("tdigest.reset", 2, "write", 1, 1, 1),
                         MakeCmdAttr<CommandTDigestMerge>("tdigest.merge", -4, "write", GetMergeKeyRange));

src/types/redis_tdigest.cc

@@ -70,6 +70,8 @@ class DummyCentroids {
       return iter_ != centroids_.cend();
     }
 
+    bool IsBegin() { return iter_ == centroids_.cbegin(); }
+
     // 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() {
@@ -186,8 +188,37 @@ rocksdb::Status TDigest::Add(engine::Context& ctx, const Slice& digest_name, con
   return storage_->Write(ctx, storage_->DefaultWriteOptions(), batch->GetWriteBatch());
 }
 
-rocksdb::Status TDigest::Quantile(engine::Context& ctx, const Slice& digest_name, const std::vector<double>& qs,
-                                  TDigestQuantitleResult* result) {
+rocksdb::Status TDigest::mergeNodes(engine::Context& ctx, const std::string& ns_key, TDigestMetadata* metadata) {
+  if (metadata->unmerged_nodes == 0) {
+    return rocksdb::Status::OK();
+  }
+
+  auto batch = storage_->GetWriteBatchBase();
+  WriteBatchLogData log_data(kRedisTDigest);
+  if (auto status = batch->PutLogData(log_data.Encode()); !status.ok()) {
+    return status;
+  }
+
+  if (auto status = mergeCurrentBuffer(ctx, ns_key, batch, metadata); !status.ok()) {
+    return status;
+  }
+
+  std::string metadata_bytes;
+  metadata->Encode(&metadata_bytes);
+  if (auto status = batch->Put(metadata_cf_handle_, ns_key, metadata_bytes); !status.ok()) {
+    return status;
+  }
+
+  if (auto status = storage_->Write(ctx, storage_->DefaultWriteOptions(), batch->GetWriteBatch()); !status.ok()) {
+    return status;
+  }
+
+  ctx.RefreshLatestSnapshot();
+  return rocksdb::Status::OK();
+}
+
+rocksdb::Status TDigest::RevRank(engine::Context& ctx, const Slice& digest_name, const std::vector<double>& inputs,
+                                 std::vector<int>& result) {
   auto ns_key = AppendNamespacePrefix(digest_name);
   TDigestMetadata metadata;
   {
@@ -198,31 +229,45 @@ rocksdb::Status TDigest::Quantile(engine::Context& ctx, const Slice& digest_name
     }
 
     if (metadata.total_observations == 0) {
+      result.resize(inputs.size(), -2);
       return rocksdb::Status::OK();
     }
 
-    if (metadata.unmerged_nodes > 0) {
-      auto batch = storage_->GetWriteBatchBase();
-      WriteBatchLogData log_data(kRedisTDigest);
-      if (auto status = batch->PutLogData(log_data.Encode()); !status.ok()) {
-        return status;
-      }
+    if (auto status = mergeNodes(ctx, ns_key, &metadata); !status.ok()) {
+      return status;
+    }
+  }
 
-      if (auto status = mergeCurrentBuffer(ctx, ns_key, batch, &metadata); !status.ok()) {
-        return status;
-      }
+  std::vector<Centroid> centroids;
+  if (auto status = dumpCentroids(ctx, ns_key, metadata, &centroids); !status.ok()) {
+    return status;
+  }
 
-      std::string metadata_bytes;
-      metadata.Encode(&metadata_bytes);
-      if (auto status = batch->Put(metadata_cf_handle_, ns_key, metadata_bytes); !status.ok()) {
-        return status;
-      }
+  auto dump_centroids = DummyCentroids(metadata, centroids);
+  auto status = TDigestRevRank(dump_centroids, inputs, result);
+  if (!status) {
+    return rocksdb::Status::InvalidArgument(status.Msg());
+  }
+  return rocksdb::Status::OK();
+}
 
-      if (auto status = storage_->Write(ctx, storage_->DefaultWriteOptions(), batch->GetWriteBatch()); !status.ok()) {
-        return status;
-      }
+rocksdb::Status TDigest::Quantile(engine::Context& ctx, const Slice& digest_name, const std::vector<double>& qs,
+                                  TDigestQuantitleResult* result) {
+  auto ns_key = AppendNamespacePrefix(digest_name);
+  TDigestMetadata metadata;
+  {
+    LockGuard guard(storage_->GetLockManager(), ns_key);
 
-      ctx.RefreshLatestSnapshot();
+    if (auto status = getMetaDataByNsKey(ctx, ns_key, &metadata); !status.ok()) {
+      return status;
+    }
+
+    if (metadata.total_observations == 0) {
+      return rocksdb::Status::OK();
+    }
+
+    if (auto status = mergeNodes(ctx, ns_key, &metadata); !status.ok()) {
+      return status;
     }
   }
 

src/types/redis_tdigest.h

@@ -77,7 +77,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 RevRank(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);
 
  private:
@@ -117,6 +118,8 @@ class TDigest : public SubKeyScanner {
   std::string internalSegmentGuardPrefixKey(const TDigestMetadata& metadata, const std::string& ns_key,
                                             SegmentType seg) const;
 
+  rocksdb::Status mergeNodes(engine::Context& ctx, const std::string& ns_key, TDigestMetadata* metadata);
+
   rocksdb::Status mergeCurrentBuffer(engine::Context& ctx, const std::string& ns_key,
                                      ObserverOrUniquePtr<rocksdb::WriteBatchBase>& batch, TDigestMetadata* metadata,
                                      const std::vector<double>* additional_buffer = nullptr,

src/types/tdigest.h

@@ -22,6 +22,8 @@
 
 #include <fmt/format.h>
 
+#include <map>
+#include <numeric>
 #include <vector>
 
 #include "common/status.h"
@@ -150,3 +152,92 @@ inline StatusOr<double> TDigestQuantile(TD&& td, double q) {
   diff /= (lc.weight / 2 + rc.weight / 2);
   return Lerp(lc.mean, rc.mean, diff);
 }
+
+inline int DoubleCompare(double a, double b, double rel_eps = 1e-12, double abs_eps = 1e-9) {
+  double diff = a - b;
+  double adiff = std::abs(diff);
+  if (adiff <= abs_eps) return 0;
+  double maxab = std::max(std::abs(a), std::abs(b));
+  if (adiff <= maxab * rel_eps) return 0;
+  return (diff < 0) ? -1 : 1;
+}
+
+inline bool DoubleEqual(double a, double b, double rel_eps = 1e-12, double abs_eps = 1e-9) {
+  return DoubleCompare(a, b, rel_eps, abs_eps) == 0;
+}
+
+struct DoubleComparator {
+  bool operator()(const double& a, const double& b) const { return DoubleCompare(a, b) == -1; }
+};
+
+template <typename TD>
+inline Status TDigestRevRank(TD&& td, const std::vector<double>& inputs, std::vector<int>& result) {
+  std::map<double, size_t, DoubleComparator> value_to_indices;
+  for (size_t i = 0; i < inputs.size(); ++i) {
+    value_to_indices[inputs[i]] = i;
+  }
+
+  result.clear();
+  result.resize(inputs.size(), -2);
+  auto it = value_to_indices.rbegin();
+
+  // handle inputs larger than maximum
+  while (it != value_to_indices.rend() && it->first > td.Max()) {
+    result[it->second] = -1;
+    ++it;
+  }
+
+  auto iter = td.End();
+  double cumulative_weight = 0;
+  while (iter->Valid() && it != value_to_indices.rend()) {
+    auto centroid = GET_OR_RET(iter->GetCentroid());
+    auto input_value = it->first;
+    if (DoubleEqual(centroid.mean, input_value)) {
+      auto current_mean = centroid.mean;
+      auto current_mean_cumulative_weight = cumulative_weight + centroid.weight / 2;
+      cumulative_weight += centroid.weight;
+
+      // handle all the previous centroids which has the same mean
+      while (!iter->IsBegin() && iter->Prev()) {
+        auto next_centroid = GET_OR_RET(iter->GetCentroid());
+        if (!DoubleEqual(current_mean, next_centroid.mean)) {
+          // move back to the last equal centroid, because we will process it in the next loop
+          iter->Next();
+          break;
+        }
+        current_mean_cumulative_weight += next_centroid.weight / 2;
+        cumulative_weight += next_centroid.weight;
+      }
+
+      // handle the prev inputs which have the same value
+      result[it->second] = static_cast<int>(current_mean_cumulative_weight);
+      ++it;
+      if (iter->IsBegin()) {
+        break;
+      }
+      iter->Prev();
+    } else if (DoubleCompare(centroid.mean, input_value) > 0) {
+      cumulative_weight += centroid.weight;
+      if (iter->IsBegin()) {
+        break;
+      }
+      iter->Prev();
+    } else {
+      result[it->second] = static_cast<int>(cumulative_weight);
+      ++it;
+    }
+  }
+
+  // handle inputs less than minimum
+  while (it != value_to_indices.rend()) {
+    result[it->second] = static_cast<int>(td.TotalWeight());
+    ++it;
+  }
+
+  for (auto r : result) {
+    if (r <= -2) {
+      return Status{Status::InvalidArgument, "invalid result when computing revrank"};
+    }
+  }
+  return Status::OK();
+}

tests/cppunit/types/tdigest_test.cc

@@ -298,3 +298,79 @@ TEST_F(RedisTDigestTest, Quantile_returns_nan_on_empty_tdigest) {
   ASSERT_TRUE(status.ok()) << status.ToString();
   ASSERT_FALSE(result.quantiles) << "should not have quantiles with empty tdigest";
 }
+
+TEST_F(RedisTDigestTest, RevRank_on_the_set_containing_different_elements) {
+  std::string test_digest_name = "test_digest_revrank" + std::to_string(util::GetTimeStampMS());
+  bool exists = false;
+  auto status = tdigest_->Create(*ctx_, test_digest_name, {100}, &exists);
+  ASSERT_FALSE(exists);
+  ASSERT_TRUE(status.ok());
+  std::vector<double> input{10, 20, 30, 40, 50, 60};
+  status = tdigest_->Add(*ctx_, test_digest_name, input);
+  ASSERT_TRUE(status.ok()) << status.ToString();
+
+  std::vector<int> result;
+  result.reserve(input.size());
+  const std::vector<double> value = {0, 10, 20, 30, 40, 50, 60, 70};
+  status = tdigest_->RevRank(*ctx_, test_digest_name, value, result);
+  const auto expect_result = std::vector<double>{6, 5, 4, 3, 2, 1, 0, -1};
+
+  for (size_t i = 0; i < result.size(); i++) {
+    auto got = result[i];
+    EXPECT_EQ(got, expect_result[i]);
+  }
+  ASSERT_TRUE(status.ok()) << status.ToString();
+}
+
+TEST_F(RedisTDigestTest, RevRank_on_the_set_containing_several_identical_elements) {
+  std::string test_digest_name = "test_digest_revrank" + std::to_string(util::GetTimeStampMS());
+  bool exists = false;
+  auto status = tdigest_->Create(*ctx_, test_digest_name, {100}, &exists);
+  ASSERT_FALSE(exists);
+  ASSERT_TRUE(status.ok());
+  std::vector<double> input{10, 10, 10, 20, 20};
+  status = tdigest_->Add(*ctx_, test_digest_name, input);
+  ASSERT_TRUE(status.ok()) << status.ToString();
+
+  std::vector<int> result;
+  result.reserve(input.size());
+  const std::vector<double> value = {10, 20};
+  status = tdigest_->RevRank(*ctx_, test_digest_name, value, result);
+  const auto expect_result = std::vector<double>{3, 1};
+  for (size_t i = 0; i < result.size(); i++) {
+    auto got = result[i];
+    EXPECT_EQ(got, expect_result[i]);
+  }
+  ASSERT_TRUE(status.ok()) << status.ToString();
+
+  status = tdigest_->Add(*ctx_, test_digest_name, std::vector<double>{10});
+  ASSERT_TRUE(status.ok()) << status.ToString();
+
+  result.clear();
+  status = tdigest_->RevRank(*ctx_, test_digest_name, value, result);
+  const auto expect_result_new = std::vector<double>{4, 1};
+  for (size_t i = 0; i < result.size(); i++) {
+    auto got = result[i];
+    EXPECT_EQ(got, expect_result_new[i]);
+  }
+  ASSERT_TRUE(status.ok()) << status.ToString();
+}
+
+TEST_F(RedisTDigestTest, RevRank_on_empty_tdigest) {
+  std::string test_digest_name = "test_digest_revrank" + std::to_string(util::GetTimeStampMS());
+  bool exists = false;
+  auto status = tdigest_->Create(*ctx_, test_digest_name, {100}, &exists);
+  ASSERT_FALSE(exists);
+  ASSERT_TRUE(status.ok());
+
+  std::vector<int> result;
+  result.reserve(2);
+  const std::vector<double> value = {10, 20};
+  status = tdigest_->RevRank(*ctx_, test_digest_name, value, result);
+  const auto expect_result = std::vector<double>{-2, -2};
+  for (size_t i = 0; i < result.size(); i++) {
+    auto got = result[i];
+    EXPECT_EQ(got, expect_result[i]);
+  }
+  ASSERT_TRUE(status.ok()) << status.ToString();
+}