feat(tdigest): implement TDIGEST.TRIMMED_MEAN command

src/commands/cmd_tdigest.cc

@@ -45,6 +45,11 @@
 constexpr auto kInfoObservations = "Observations";
 constexpr auto kInfoTotalCompressions = "Total compressions";
 constexpr auto kNan = "nan";
+
+constexpr const char *errParseLowCutQuantile = "error parsing low_cut_percentile";
+constexpr const char *errParseHighCutQuantile = "error parsing high_cut_percentile";
+constexpr const char *errCutQuantileRange = "low_cut_percentile and high_cut_percentile should be in [0,1]";
+constexpr const char *errLowCutQuantileLess = "low_cut_percentile should be lower than high_cut_percentile";
 }  // namespace
 
 class CommandTDigestCreate : public Commander {
@@ -492,6 +497,67 @@
   TDigestMergeOptions options_;
 };
 
+class CommandTDigestTrimmedMean : public Commander {
+ public:
+  Status Parse(const std::vector<std::string> &args) override {
+    if (args.size() != 4) {
+      return {Status::RedisParseErr, errWrongNumOfArguments};
+    }
+
+    key_name_ = args[1];
+
+    auto low_cut_quantile = ParseFloat(args[2]);
+    if (!low_cut_quantile || std::isnan(*low_cut_quantile)) {
+      return {Status::RedisParseErr, errParseLowCutQuantile};
+    }
+    low_cut_quantile_ = *low_cut_quantile;
+
+    auto high_cut_quantile = ParseFloat(args[3]);
+    if (!high_cut_quantile || std::isnan(*high_cut_quantile)) {
+      return {Status::RedisParseErr, errParseHighCutQuantile};
+    }
+    high_cut_quantile_ = *high_cut_quantile;
+
+    if (!std::isfinite(low_cut_quantile_) || low_cut_quantile_ < 0.0 || low_cut_quantile_ > 1.0) {
+      return {Status::RedisParseErr, errCutQuantileRange};
+    }
+    if (!std::isfinite(high_cut_quantile_) || high_cut_quantile_ < 0.0 || high_cut_quantile_ > 1.0) {
+      return {Status::RedisParseErr, errCutQuantileRange};
+    }
+    if (low_cut_quantile_ >= high_cut_quantile_) {
+      return {Status::RedisParseErr, errLowCutQuantileLess};
+    }
+
+    return Status::OK();
+  }
+
+  Status Execute(engine::Context &ctx, Server *srv, Connection *conn, std::string *output) override {
+    TDigest tdigest(srv->storage, conn->GetNamespace());
+    TDigestTrimmedMeanResult result;
+
+    auto s = tdigest.TrimmedMean(ctx, key_name_, low_cut_quantile_, high_cut_quantile_, &result);
+    if (!s.ok()) {
+      if (s.IsNotFound()) {
+        return {Status::RedisExecErr, errKeyNotFound};
+      }
+      return {Status::RedisExecErr, s.ToString()};
+    }
+
+    if (!result.mean.has_value()) {
+      *output = redis::BulkString(kNan);
+    } else {
+      *output = redis::BulkString(util::Float2String(*result.mean));
+    }
+
+    return Status::OK();
+  }
+
+ private:
+  std::string key_name_;
+  double low_cut_quantile_;
+  double high_cut_quantile_;
+};
+
 std::vector<CommandKeyRange> GetMergeKeyRange(const std::vector<std::string> &args) {
   auto numkeys = ParseInt<int>(args[2], 10).ValueOr(0);
   return {{1, 1, 1}, {3, 2 + numkeys, 1}};
@@ -507,6 +573,7 @@
                         MakeCmdAttr<CommandTDigestByRevRank>("tdigest.byrevrank", -3, "read-only", 1, 1, 1),
                         MakeCmdAttr<CommandTDigestByRank>("tdigest.byrank", -3, "read-only", 1, 1, 1),
                         MakeCmdAttr<CommandTDigestQuantile>("tdigest.quantile", -3, "read-only", 1, 1, 1),
+                        MakeCmdAttr<CommandTDigestTrimmedMean>("tdigest.trimmed_mean", 4, "read-only", 1, 1, 1),
                         MakeCmdAttr<CommandTDigestReset>("tdigest.reset", 2, "write", 1, 1, 1),
                         MakeCmdAttr<CommandTDigestMerge>("tdigest.merge", -4, "write", GetMergeKeyRange));
 }  // namespace redis

src/types/redis_tdigest.cc

@@ -759,6 +759,42 @@ rocksdb::Status TDigest::applyNewCentroids(ObserverOrUniquePtr<rocksdb::WriteBat
   return rocksdb::Status::OK();
 }
 
+rocksdb::Status TDigest::TrimmedMean(engine::Context& ctx, const Slice& digest_name, double low_cut_quantile,
+                                     double high_cut_quantile, TDigestTrimmedMeanResult* 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->mean.reset();
+      return rocksdb::Status::OK();
+    }
+
+    if (auto status = mergeNodes(ctx, ns_key, &metadata); !status.ok()) {
+      return status;
+    }
+  }
+
+  // Dump centroids and create DummyCentroids wrapper for TDigest algorithm
+  std::vector<Centroid> centroids;
+  if (auto status = dumpCentroids(ctx, ns_key, metadata, &centroids); !status.ok()) {
+    return status;
+  }
+  auto dump_centroids = DummyCentroids<false>(metadata, centroids);
+  auto trimmed_mean_result = TDigestTrimmedMean(dump_centroids, low_cut_quantile, high_cut_quantile);
+  if (!trimmed_mean_result) {
+    return rocksdb::Status::InvalidArgument(trimmed_mean_result.Msg());
+  }
+
+  result->mean = *trimmed_mean_result;
+  return rocksdb::Status::OK();
+}
+
 std::string TDigest::internalSegmentGuardPrefixKey(const TDigestMetadata& metadata, const std::string& ns_key,
                                                    SegmentType seg) const {
   std::string prefix_key;

src/types/redis_tdigest.h

@@ -53,6 +53,10 @@ struct TDigestQuantitleResult {
   std::optional<std::vector<double>> quantiles;
 };
 
+struct TDigestTrimmedMeanResult {
+  std::optional<double> mean;
+};
+
 class TDigest : public SubKeyScanner {
  public:
   using Slice = rocksdb::Slice;
@@ -85,6 +89,8 @@ class TDigest : public SubKeyScanner {
                             std::vector<double>* result);
   rocksdb::Status ByRank(engine::Context& ctx, const Slice& digest_name, const std::vector<int>& inputs,
                          std::vector<double>* result);
+  rocksdb::Status TrimmedMean(engine::Context& ctx, const Slice& digest_name, double low_cut_quantile,
+                              double high_cut_quantile, TDigestTrimmedMeanResult* result);
   rocksdb::Status GetMetaData(engine::Context& context, const Slice& digest_name, TDigestMetadata* metadata);
 
  private:

src/types/tdigest.h

@@ -309,3 +309,46 @@
   }
   return Status::OK();
 }
+
+template <typename TD>
+inline StatusOr<double> TDigestTrimmedMean(TD&& td, double low_cut_quantile, double high_cut_quantile) {
+  if (td.Size() == 0) {
+    return std::numeric_limits<double>::quiet_NaN();
+  }
+
+  const double total_weight = td.TotalWeight();
+  const double leftmost_weight = std::floor(total_weight * low_cut_quantile);
+  const double rightmost_weight = std::ceil(total_weight * high_cut_quantile);
+
+  double count_done = 0.0;
+  double trimmed_sum = 0.0;
+  double trimmed_count = 0.0;
+
+  auto iter = td.Begin();
+  while (iter->Valid()) {
+    auto centroid = GET_OR_RET(iter->GetCentroid());
+    const double n_weight = centroid.weight;
+    double count_add = n_weight;
+
+    // Keep only the portion of this centroid that overlaps with the trimmed rank range.
+    count_add -= std::min(std::max(0.0, leftmost_weight - count_done), count_add);
+    count_add = std::min(std::max(0.0, rightmost_weight - count_done), count_add);
+
+    count_done += n_weight;
+
+    trimmed_sum += centroid.mean * count_add;
+    trimmed_count += count_add;
+
+    if (count_done >= rightmost_weight) {
+      break;
+    }
+
+    iter->Next();
+  }
+
+  if (trimmed_count == 0.0) {
+    return std::numeric_limits<double>::quiet_NaN();
+  }
+
+  return trimmed_sum / trimmed_count;
+}

tests/cppunit/types/tdigest_test.cc

@@ -524,3 +524,81 @@ TEST_F(RedisTDigestTest, ByRank_And_ByRevRank) {
   EXPECT_EQ(result[0], 1.0) << "Rank 0 should be minimum";
   EXPECT_TRUE(std::isinf(result[3])) << "Rank >= total_weight should be infinity";
 }
+
+TEST_F(RedisTDigestTest, TrimmedMean) {
+  std::string test_digest_name = "test_digest_trimmed_mean" + 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> values = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10};
+  status = tdigest_->Add(*ctx_, test_digest_name, values);
+  ASSERT_TRUE(status.ok()) << status.ToString();
+
+  redis::TDigestTrimmedMeanResult result;
+  status = tdigest_->TrimmedMean(*ctx_, test_digest_name, 0.1, 0.9, &result);
+  ASSERT_TRUE(status.ok()) << status.ToString();
+  ASSERT_TRUE(result.mean.has_value());
+  EXPECT_NEAR(*result.mean, 5.5, 0.01);
+
+  status = tdigest_->TrimmedMean(*ctx_, test_digest_name, 0.0, 1.0, &result);
+  ASSERT_TRUE(status.ok()) << status.ToString();
+  ASSERT_TRUE(result.mean.has_value());
+  EXPECT_NEAR(*result.mean, 5.5, 0.01);
+
+  status = tdigest_->TrimmedMean(*ctx_, test_digest_name, 0.25, 0.75, &result);
+  ASSERT_TRUE(status.ok()) << status.ToString();
+  ASSERT_TRUE(result.mean.has_value());
+  EXPECT_NEAR(*result.mean, 5.5, 0.01);
+}
+
+TEST_F(RedisTDigestTest, TrimmedMeanEmptyDigest) {
+  std::string test_digest_name = "test_digest_trimmed_mean_empty" + 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());
+
+  redis::TDigestTrimmedMeanResult result;
+  status = tdigest_->TrimmedMean(*ctx_, test_digest_name, 0.1, 0.9, &result);
+  ASSERT_TRUE(status.ok()) << status.ToString();
+  ASSERT_FALSE(result.mean.has_value());
+}
+
+TEST_F(RedisTDigestTest, TrimmedMeanUnorderedInput) {
+  std::string test_digest_name = "test_digest_trimmed_mean_unordered" + 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> values = {5, 2, 8, 1, 9, 3, 7, 4, 6, 10};
+  status = tdigest_->Add(*ctx_, test_digest_name, values);
+  ASSERT_TRUE(status.ok()) << status.ToString();
+
+  redis::TDigestTrimmedMeanResult result;
+  status = tdigest_->TrimmedMean(*ctx_, test_digest_name, 0.1, 0.9, &result);
+  ASSERT_TRUE(status.ok()) << status.ToString();
+  ASSERT_TRUE(result.mean.has_value());
+  EXPECT_NEAR(*result.mean, 5.5, 0.01);
+}
+
+TEST_F(RedisTDigestTest, TrimmedMeanComplexInput) {
+  std::string test_digest_name = "test_digest_trimmed_mean_complex" + 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> values = {-10, 5, -3, 5, 0, 5, 3, -5, 10, -10};
+  status = tdigest_->Add(*ctx_, test_digest_name, values);
+  ASSERT_TRUE(status.ok()) << status.ToString();
+
+  redis::TDigestTrimmedMeanResult result;
+  status = tdigest_->TrimmedMean(*ctx_, test_digest_name, 0.2, 0.8, &result);
+  ASSERT_TRUE(status.ok()) << status.ToString();
+  ASSERT_TRUE(result.mean.has_value());
+  ASSERT_FALSE(std::isnan(*result.mean));
+  EXPECT_NEAR(*result.mean, 5.0 / 6.0, 0.01);
+}