feat(cms): add Count-Min Sketch data structure

Tangruilin · Tangruilin · commit ce091bf39c3e · 2026-03-29T19:49:26.000+08:00
Implement CMS (Count-Min Sketch) probabilistic data structure with the following commands: - CMS.INITBYDIM key width depth - CMS.INITBYPROB key error probability - CMS.INCRBY key item increment [...] - CMS.QUERY key item [...] - CMS.MERGE dest numkeys src... [WEIGHTS ...] - CMS.INFO key Storage design: - Metadata stored in Metadata CF - Count Matrix stored in PrimarySubkey CF (per-bucket storage) - Uses MurmurHash64 for layer-specific hashing Related: #2425
diff --git a/src/commands/cmd_cms.cc b/src/commands/cmd_cms.cc
@@ -0,0 +1,272 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ *
+ */
+
+#include <types/redis_cms.h>
+
+#include "commander.h"
+#include "commands/command_parser.h"
+#include "server/redis_reply.h"
+#include "server/server.h"
+
+namespace redis {
+
+/// CMS.INITBYDIM key width depth
+/// Initialize a Count-Min Sketch with specified dimensions.
+/// Complexity: O(width * depth) to initialize all buckets.
+class CommandCMSInitByDim final : public Commander {
+ public:
+  Status Parse(const std::vector<std::string> &args) override {
+    auto parse_width = ParseInt<uint32_t>(args[2], 10);
+    if (!parse_width) {
+      return {Status::RedisParseErr, "invalid width"};
+    }
+    width_ = *parse_width;
+
+    auto parse_depth = ParseInt<uint32_t>(args[3], 10);
+    if (!parse_depth) {
+      return {Status::RedisParseErr, "invalid depth"};
+    }
+    depth_ = *parse_depth;
+
+    return Commander::Parse(args);
+  }
+
+  Status Execute(engine::Context &ctx, Server *srv, Connection *conn, std::string *output) override {
+    redis::CMS cms(srv->storage, conn->GetNamespace());
+
+    auto s = cms.InitByDim(ctx, args_[1], width_, depth_);
+    if (!s.ok()) return {Status::RedisExecErr, s.ToString()};
+
+    *output = redis::RESP_OK;
+    return Status::OK();
+  }
+
+ private:
+  uint32_t width_;
+  uint32_t depth_;
+};
+
+/// CMS.INITBYPROB key error probability
+/// Initialize a Count-Min Sketch with specified error rate and probability.
+/// Complexity: O(width * depth) where width/depth are calculated from error/probability.
+class CommandCMSInitByProb final : public Commander {
+ public:
+  Status Parse(const std::vector<std::string> &args) override {
+    auto parse_error = ParseFloat<double>(args[2]);
+    if (!parse_error) {
+      return {Status::RedisParseErr, "invalid error rate"};
+    }
+    error_rate_ = *parse_error;
+
+    auto parse_prob = ParseFloat<double>(args[3]);
+    if (!parse_prob) {
+      return {Status::RedisParseErr, "invalid probability"};
+    }
+    probability_ = *parse_prob;
+
+    return Commander::Parse(args);
+  }
+
+  Status Execute(engine::Context &ctx, Server *srv, Connection *conn, std::string *output) override {
+    redis::CMS cms(srv->storage, conn->GetNamespace());
+
+    auto s = cms.InitByProb(ctx, args_[1], error_rate_, probability_);
+    if (!s.ok()) return {Status::RedisExecErr, s.ToString()};
+
+    *output = redis::RESP_OK;
+    return Status::OK();
+  }
+
+ private:
+  double error_rate_;
+  double probability_;
+};
+
+/// CMS.INCRBY key item increment [item increment ...]
+/// Increment the count of one or more items.
+/// Complexity: O(depth) for each item.
+/// Returns: Array of estimated counts for each item after increment.
+class CommandCMSIncrBy final : public Commander {
+ public:
+  Status Parse(const std::vector<std::string> &args) override {
+    if (args.size() < 4 || (args.size() - 2) % 2 != 0) {
+      return {Status::RedisParseErr, "wrong number of arguments"};
+    }
+
+    for (size_t i = 2; i < args.size(); i += 2) {
+      auto parse_increment = ParseInt<int64_t>(args[i + 1], 10);
+      if (!parse_increment) {
+        return {Status::RedisParseErr, "invalid increment"};
+      }
+      if (*parse_increment < 0) {
+        return {Status::RedisParseErr, "increment must be non-negative"};
+      }
+      items_.emplace_back(args[i], *parse_increment);
+    }
+
+    return Commander::Parse(args);
+  }
+
+  Status Execute(engine::Context &ctx, Server *srv, Connection *conn, std::string *output) override {
+    redis::CMS cms(srv->storage, conn->GetNamespace());
+    std::vector<uint64_t> counts;
+
+    auto s = cms.IncrBy(ctx, args_[1], items_, &counts);
+    if (!s.ok()) return {Status::RedisExecErr, s.ToString()};
+
+    *output = redis::MultiLen(counts.size());
+    for (auto count : counts) {
+      *output += redis::Integer(count);
+    }
+    return Status::OK();
+  }
+
+ private:
+  std::vector<std::pair<std::string, int64_t>> items_;
+};
+
+/// CMS.QUERY key item [item ...]
+/// Return the estimated count of one or more items.
+/// Complexity: O(depth) for each item.
+/// Returns: Array of estimated counts for each item.
+class CommandCMSQuery final : public Commander {
+ public:
+  Status Parse(const std::vector<std::string> &args) override {
+    items_.reserve(args.size() - 2);
+    for (size_t i = 2; i < args.size(); ++i) {
+      items_.push_back(args[i]);
+    }
+    return Commander::Parse(args);
+  }
+
+  Status Execute(engine::Context &ctx, Server *srv, Connection *conn, std::string *output) override {
+    redis::CMS cms(srv->storage, conn->GetNamespace());
+    std::vector<uint64_t> counts;
+
+    auto s = cms.Query(ctx, args_[1], items_, &counts);
+    if (!s.ok()) return {Status::RedisExecErr, s.ToString()};
+
+    *output = redis::MultiLen(counts.size());
+    for (auto count : counts) {
+      *output += redis::Integer(count);
+    }
+    return Status::OK();
+  }
+
+ private:
+  std::vector<std::string> items_;
+};
+
+/// CMS.MERGE destkey numkeys sourcekey [sourcekey ...] [WEIGHTS weight [weight ...]]
+/// Merge multiple Count-Min Sketches into one.
+/// Complexity: O(width * depth * numkeys).
+///
+/// Kvrocks extension:
+///   The underlying implementation supports SUM/MAX/MIN merge methods via CMSMergeMethod enum.
+///   Currently only SUM is exposed for Redis compatibility. Future versions may add:
+///   [METHOD SUM|MAX|MIN] parameter to allow specifying merge strategy.
+///   This makes Kvrocks a superset of Redis CMS functionality.
+class CommandCMSMerge final : public Commander {
+ public:
+  Status Parse(const std::vector<std::string> &args) override {
+    auto parse_numkeys = ParseInt<size_t>(args[2], 10);
+    if (!parse_numkeys) {
+      return {Status::RedisParseErr, "invalid numkeys"};
+    }
+    numkeys_ = *parse_numkeys;
+
+    if (args.size() < 3 + numkeys_) {
+      return {Status::RedisParseErr, "wrong number of arguments"};
+    }
+
+    // Parse source keys
+    for (size_t i = 0; i < numkeys_; ++i) {
+      src_keys_.push_back(args[3 + i]);
+    }
+
+    // Parse optional WEIGHTS
+    size_t next_arg = 3 + numkeys_;
+    if (next_arg < args.size() && strcasecmp(args[next_arg].c_str(), "WEIGHTS") == 0) {
+      next_arg++;
+      if (args.size() < next_arg + numkeys_) {
+        return {Status::RedisParseErr, "wrong number of weights"};
+      }
+      for (size_t i = 0; i < numkeys_; ++i) {
+        auto parse_weight = ParseInt<uint64_t>(args[next_arg + i], 10);
+        if (!parse_weight) {
+          return {Status::RedisParseErr, "invalid weight"};
+        }
+        weights_.push_back(*parse_weight);
+      }
+    }
+
+    return Commander::Parse(args);
+  }
+
+  Status Execute(engine::Context &ctx, Server *srv, Connection *conn, std::string *output) override {
+    redis::CMS cms(srv->storage, conn->GetNamespace());
+
+    auto s = cms.Merge(ctx, args_[1], src_keys_, weights_, CMSMergeMethod::SUM);
+    if (!s.ok()) return {Status::RedisExecErr, s.ToString()};
+
+    *output = redis::RESP_OK;
+    return Status::OK();
+  }
+
+ private:
+  size_t numkeys_;
+  std::vector<std::string> src_keys_;
+  std::vector<uint64_t> weights_;
+};
+
+/// CMS.INFO key
+/// Return information about a Count-Min Sketch.
+/// Returns: Array of key-value pairs (width, depth, total count, size).
+class CommandCMSInfo final : public Commander {
+ public:
+  Status Execute(engine::Context &ctx, Server *srv, Connection *conn, std::string *output) override {
+    redis::CMS cms(srv->storage, conn->GetNamespace());
+    CMSInfo info;
+
+    auto s = cms.Info(ctx, args_[1], &info);
+    if (s.IsNotFound()) return {Status::RedisExecErr, "key not found"};
+    if (!s.ok()) return {Status::RedisExecErr, s.ToString()};
+
+    *output = redis::MultiLen(8);
+    *output += redis::SimpleString("width");
+    *output += redis::Integer(info.width);
+    *output += redis::SimpleString("depth");
+    *output += redis::Integer(info.depth);
+    *output += redis::SimpleString("count");
+    *output += redis::Integer(info.total_count);
+    *output += redis::SimpleString("size");
+    *output += redis::Integer(info.size);
+    return Status::OK();
+  }
+};
+
+REDIS_REGISTER_COMMANDS(CMS, MakeCmdAttr<CommandCMSInitByDim>("cms.initbydim", 4, "write", 1, 1, 1),
+                        MakeCmdAttr<CommandCMSInitByProb>("cms.initbyprob", 4, "write", 1, 1, 1),
+                        MakeCmdAttr<CommandCMSIncrBy>("cms.incrby", -4, "write", 1, 1, 1),
+                        MakeCmdAttr<CommandCMSQuery>("cms.query", -3, "read-only", 1, 1, 1),
+                        MakeCmdAttr<CommandCMSMerge>("cms.merge", -4, "write", 1, 1, 1),
+                        MakeCmdAttr<CommandCMSInfo>("cms.info", 2, "read-only", 1, 1, 1), )
+
+}  // namespace redis
diff --git a/src/commands/commander.h b/src/commands/commander.h
@@ -95,6 +95,7 @@ enum class CommandCategory : uint8_t {
   Unknown = 0,
   Bit,
   BloomFilter,
+  CMS,
   Cluster,
   Function,
   Geo,
diff --git a/src/storage/redis_metadata.cc b/src/storage/redis_metadata.cc
@@ -334,7 +334,7 @@ bool Metadata::IsSingleKVType() const { return Type() == kRedisString || Type()
 
 bool Metadata::IsEmptyableType() const {
   return IsSingleKVType() || Type() == kRedisStream || Type() == kRedisBloomFilter || Type() == kRedisHyperLogLog ||
-         Type() == kRedisTDigest || Type() == kRedisTimeSeries;
+         Type() == kRedisTDigest || Type() == kRedisTimeSeries || Type() == kRedisCMS;
 }
 
 bool Metadata::Expired() const { return ExpireAt(util::GetTimeStampMS()); }
@@ -569,3 +569,32 @@ rocksdb::Status TimeSeriesMetadata::Decode(Slice *input) {
 
   return rocksdb::Status::OK();
 }
+
+void CMSMetadata::Encode(std::string *dst) const {
+  Metadata::Encode(dst);
+
+  PutFixed32(dst, width);
+  PutFixed32(dst, depth);
+  PutFixed64(dst, total_count);
+  PutFixed8(dst, static_cast<uint8_t>(storage_mode));
+}
+
+rocksdb::Status CMSMetadata::Decode(Slice *input) {
+  if (auto s = Metadata::Decode(input); !s.ok()) {
+    return s;
+  }
+
+  if (input->size() < 4 + 4 + 8 + 1) {
+    return rocksdb::Status::InvalidArgument(kErrMetadataTooShort);
+  }
+
+  GetFixed32(input, &width);
+  GetFixed32(input, &depth);
+  GetFixed64(input, &total_count);
+
+  uint8_t mode = 0;
+  GetFixed8(input, &mode);
+  storage_mode = static_cast<StorageMode>(mode);
+
+  return rocksdb::Status::OK();
+}
diff --git a/src/storage/redis_metadata.h b/src/storage/redis_metadata.h
@@ -54,12 +54,13 @@ enum RedisType : uint8_t {
   kRedisHyperLogLog = 11,
   kRedisTDigest = 12,
   kRedisTimeSeries = 13,
+  kRedisCMS = 14,
   kRedisTypeMax
 };
 
 inline constexpr const std::array<std::string_view, kRedisTypeMax> RedisTypeNames = {
     "none",      "string", "hash",      "list",      "set",         "zset",      "bitmap",
-    "sortedint", "stream", "MBbloom--", "ReJSON-RL", "hyperloglog", "TDIS-TYPE", "timeseries"};
+    "sortedint", "stream", "MBbloom--", "ReJSON-RL", "hyperloglog", "TDIS-TYPE", "timeseries", "cms"};
 
 struct RedisTypes {
   RedisTypes(std::initializer_list<RedisType> list) {
@@ -409,3 +410,32 @@ class TimeSeriesMetadata : public Metadata {
   void Encode(std::string *dst) const override;
   rocksdb::Status Decode(Slice *input) override;
 };
+
+class CMSMetadata : public Metadata {
+ public:
+  enum class StorageMode : uint8_t {
+    PER_BUCKET = 0,    // 按桶存储（默认）
+    SINGLE_KEY = 1,    // 单 Key 存储
+  };
+
+  /// Width of the count matrix (number of buckets per layer)
+  uint32_t width;
+
+  /// Depth of the count matrix (number of layers)
+  uint32_t depth;
+
+  /// Total count of all INCRBY operations
+  uint64_t total_count;
+
+  /// Storage mode
+  StorageMode storage_mode;
+
+  explicit CMSMetadata(bool generate_version = true)
+      : Metadata(kRedisCMS, generate_version), width(0), depth(0), total_count(0), storage_mode(StorageMode::PER_BUCKET) {}
+
+  CMSMetadata(uint32_t width, uint32_t depth, bool generate_version = true)
+      : Metadata(kRedisCMS, generate_version), width(width), depth(depth), total_count(0), storage_mode(StorageMode::PER_BUCKET) {}
+
+  void Encode(std::string *dst) const override;
+  rocksdb::Status Decode(Slice *input) override;
+};
diff --git a/src/types/redis_cms.cc b/src/types/redis_cms.cc
diff --git a/src/types/redis_cms.h b/src/types/redis_cms.h
