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

src/commands/cmd_tdigest.cc

@@ -176,6 +176,49 @@ 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];
+    inputs_.reserve(args.size() - 2);
+    for (size_t i = 2; i < args.size(); i++) {
+      auto value = ParseFloat(args[i]);
+      if (!value) {
+        return {Status::RedisParseErr, errValueIsNotFloat};
+      }
+      inputs_.push_back(*value);
+    }
+    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(inputs_.size());
+    if (const auto s = tdigest.RevRank(ctx, key_name_, inputs_, result); !s.ok()) {
+      if (s.IsNotFound()) {
+        return {Status::RedisExecErr, errKeyNotFound};
+      }
+      return {Status::RedisExecErr, s.ToString()};
+    }
+
+    if (!result.empty()) {
+      std::vector<std::string> rev_ranks;
+      rev_ranks.reserve(result.size());
+      for (const auto v : result) {
+        rev_ranks.push_back(redis::Integer(v));
+      }
+      *output = redis::Array(rev_ranks);
+    } else {
+      *output = redis::BulkString("nan");
+    }
+    return Status::OK();
+  }
+
+ private:
+  std::string key_name_;
+  std::vector<double> inputs_;
+};
+
 class CommandTDigestMinMax : public Commander {
  public:
   explicit CommandTDigestMinMax(bool is_min) : is_min_(is_min) {}
@@ -369,6 +412,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

@@ -67,18 +67,18 @@ class DummyCentroids {
       if (Valid()) {
         std::advance(iter_, 1);
       }
-      return iter_ != centroids_.cend();
+      return Valid();
     }
 
     // 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_ != centroids_.cbegin()) {
+      if (Valid()) {
         std::advance(iter_, -1);
       }
       return Valid();
     }
-    bool Valid() const { return iter_ != centroids_.cend(); }
+    bool Valid() const { return iter_ < centroids_.cend() && iter_ >= centroids_.cbegin(); }
     StatusOr<Centroid> GetCentroid() const {
       if (iter_ == centroids_.cend()) {
         return {::Status::NotOK, "invalid iterator during decoding tdigest centroid"};
@@ -186,8 +186,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 +227,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,70 @@ inline StatusOr<double> TDigestQuantile(TD&& td, double q) {
   diff /= (lc.weight / 2 + rc.weight / 2);
   return Lerp(lc.mean, rc.mean, diff);
 }
+
+inline void AssignRankForEqualInputs(const std::vector<size_t>& indices, double cumulative_weight,
+                                     std::vector<int>& result) {
+  for (auto index : indices) {
+    result[index] = static_cast<int>(cumulative_weight);
+  }
+}
+
+template <typename TD>
+inline Status TDigestRevRank(TD&& td, const std::vector<double>& inputs, std::vector<int>& result) {
+  std::map<double, std::vector<size_t>> value_to_indices;
+  for (size_t i = 0; i < inputs.size(); ++i) {
+    value_to_indices[inputs[i]].push_back(i);
+  }
+
+  double cumulative_weight = 0;
+  result.resize(inputs.size());
+  auto it = value_to_indices.rbegin();
+
+  // handle inputs larger than maximum
+  while (it != value_to_indices.rend() && it->first > td.Max()) {
+    AssignRankForEqualInputs(it->second, -1, result);
+    ++it;
+  }
+
+  auto iter = td.End();
+  while (iter->Valid() && it != value_to_indices.rend()) {
+    auto centroid = GET_OR_RET(iter->GetCentroid());
+    auto input_value = it->first;
+    if (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 prev centroids which has the same mean
+      while (iter->Prev()) {
+        auto next_centroid = GET_OR_RET(iter->GetCentroid());
+        if (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 += centroid.weight / 2;
+        cumulative_weight += centroid.weight;
+      }
+
+      // handle the prev inputs which has the same value
+      AssignRankForEqualInputs(it->second, current_mean_cumulative_weight, result);
+      ++it;
+      iter->Prev();
+    } else if (centroid.mean > input_value) {
+      cumulative_weight += centroid.weight;
+      iter->Prev();
+    } else {
+      AssignRankForEqualInputs(it->second, cumulative_weight, result);
+      ++it;
+    }
+  }
+
+  // handle inputs less than minimum
+  while (it != value_to_indices.rend()) {
+    AssignRankForEqualInputs(it->second, td.TotalWeight(), result);
+    ++it;
+  }
+
+  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_contains_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_contains_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();
+}