feat: support count min sketch data structure and most commands

src/commands/cmd_cms.cc

@@ -0,0 +1,168 @@
+/*
+ * 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/cms.h>
+#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.INCRBY key item increment [item increment ...]
+class CommandCMSIncrBy final : public Commander {
+ public:
+  Status Execute(Server *srv, Connection *conn, std::string *output) override {
+    if ((args_.size() - 2) % 2 != 0) {
+      return Status::RedisTryAgain;
+    }
+    redis::CMS cms(srv->storage, conn->GetNamespace());
+    engine::Context ctx(srv->storage);
+    rocksdb::Status s;
+    std::unordered_map<std::string, uint64_t> elements;
+    for (size_t i = 2; i < args_.size(); i += 2) {
+      std::string key = args_[i];
+      uint64_t value = 0;
+      try {
+        value = std::stoull(args_[i + 1]);
+      } catch (const std::exception &e) {
+        return Status::InvalidArgument;
+      }
+      elements[key] = value;
+    }
+
+    s = cms.IncrBy(ctx, args_[1], elements);
+    if (!s.ok()) {
+      return {Status::RedisExecErr, s.ToString()};
+    }
+
+    *output = redis::SimpleString("OK");
+    return Status::OK();
+  }
+};
+
+/// CMS.INFO key
+class CommandCMSInfo final : public Commander {
+ public:
+  Status Execute(Server *srv, Connection *conn, std::string *output) override {
+    redis::CMS cms(srv->storage, conn->GetNamespace());
+    engine::Context ctx(srv->storage);
+    rocksdb::Status s;
+    std::unordered_map<std::string, uint64_t> elements;
+    std::vector<uint64_t> ret{};
+
+    s = cms.Info(ctx, args_[1], &ret);
+
+    if (s.IsNotFound()) {
+      return {Status::RedisExecErr, s.ToString()};
+    }
+
+    if (!s.ok() && !s.IsNotFound()) {
+      return {Status::RedisExecErr, s.ToString()};
+    }
+
+    *output = redis::Array({redis::BulkString("width"), redis::Integer(ret[0]), redis::BulkString("depth"),
+                            redis::Integer(ret[1]), redis::BulkString("count"), redis::Integer(ret[2])});
+
+    return Status::OK();
+  }
+};
+
+/// CMS.INITBYDIM key width depth
+class CommandCMSInitByDim final : public Commander {
+ public:
+  Status Execute(Server *srv, Connection *conn, std::string *output) override {
+    redis::CMS cms(srv->storage, conn->GetNamespace());
+    engine::Context ctx(srv->storage);
+    rocksdb::Status s;
+    uint64_t width = std::stoull(args_[2]);
+    uint64_t depth = std::stoull(args_[3]);
+
+    s = cms.InitByDim(ctx, args_[1], width, depth);
+    if (!s.ok()) {
+      return {Status::RedisExecErr, s.ToString()};
+    }
+
+    *output = redis::SimpleString("OK");
+    return Status::OK();
+  }
+};
+
+/// CMS.INITBYPROB key error probability
+class CommandCMSInitByProb final : public Commander {
+ public:
+  Status Execute(Server *srv, Connection *conn, std::string *output) override {
+    redis::CMS cms(srv->storage, conn->GetNamespace());
+    engine::Context ctx(srv->storage);
+    rocksdb::Status s;
+    double error = std::stod(args_[2]);
+    double delta = std::stod(args_[3]);
+
+    s = cms.InitByProb(ctx, args_[1], error, delta);
+    if (!s.ok()) {
+      return {Status::RedisExecErr, s.ToString()};
+    }
+
+    *output = redis::SimpleString("OK");
+    return Status::OK();
+  }
+};
+
+/// CMS.QUERY key item [item ...]
+class CommandCMSQuery final : public Commander {
+ public:
+  Status Execute(Server *srv, Connection *conn, std::string *output) override {
+    redis::CMS cms(srv->storage, conn->GetNamespace());
+    engine::Context ctx(srv->storage);
+    rocksdb::Status s;
+
+    std::vector<uint32_t> counters{};
+    std::vector<std::string> elements;
+
+    for (size_t i = 2; i < args_.size(); ++i) {
+      elements.emplace_back(args_[i]);
+    }
+
+    s = cms.Query(ctx, args_[1], elements, counters);
+
+    if (!s.ok()) {
+      return {Status::RedisExecErr, s.ToString()};
+    }
+
+    std::vector<std::string> output_values;
+    output_values.reserve(counters.size());
+    for (const auto &counter : counters) {
+      output_values.push_back(std::to_string(counter));
+    }
+
+    *output = redis::ArrayOfBulkStrings(output_values);
+
+    return Status::OK();
+  }
+};
+
+REDIS_REGISTER_COMMANDS(CMS, MakeCmdAttr<CommandCMSIncrBy>("cms.incrby", -4, "write", 0, 0, 0),
+                        MakeCmdAttr<CommandCMSInfo>("cms.info", 2, "read-only", 0, 0, 0),
+                        MakeCmdAttr<CommandCMSInitByDim>("cms.initbydim", 4, "write", 0, 0, 0),
+                        MakeCmdAttr<CommandCMSInitByProb>("cms.initbyprob", 4, "write", 0, 0, 0),
+                        MakeCmdAttr<CommandCMSQuery>("cms.query", -3, "read-only", 0, 0, 0), );
+}  // namespace redis

src/commands/command_parser.h

@@ -33,7 +33,7 @@
 
 template <typename Iter>
 struct MoveIterator : Iter {
-  explicit MoveIterator(Iter iter) : Iter(iter){};
+  explicit MoveIterator(Iter iter) : Iter(iter) {};
 
   typename Iter::value_type&& operator*() const { return std::move(this->Iter::operator*()); }
 };

src/commands/commander.h

@@ -73,6 +73,7 @@ enum class CommandCategory : uint8_t {
   Bit,
   BloomFilter,
   Cluster,
+  CMS,
   Function,
   Geo,
   Hash,

src/storage/redis_metadata.cc

@@ -329,7 +329,8 @@ bool Metadata::ExpireAt(uint64_t expired_ts) const {
 bool Metadata::IsSingleKVType() const { return Type() == kRedisString || Type() == kRedisJson; }
 
 bool Metadata::IsEmptyableType() const {
-  return IsSingleKVType() || Type() == kRedisStream || Type() == kRedisBloomFilter || Type() == kRedisHyperLogLog;
+  return IsSingleKVType() || Type() == kRedisStream || Type() == kRedisBloomFilter || Type() == kRedisHyperLogLog ||
+         Type() == kRedisCountMinSketch;
 }
 
 bool Metadata::Expired() const { return ExpireAt(util::GetTimeStampMS()); }
@@ -495,3 +496,41 @@ rocksdb::Status HyperLogLogMetadata::Decode(Slice *input) {
 
   return rocksdb::Status::OK();
 }
+
+void CountMinSketchMetadata::Encode(std::string *dst) const {
+  Metadata::Encode(dst);
+  PutFixed32(dst, width);
+  PutFixed32(dst, depth);
+  PutFixed64(dst, counter);
+  for (const auto &count : array) {
+    PutFixed32(dst, count);
+  }
+}
+
+rocksdb::Status CountMinSketchMetadata::Decode(Slice *input) {
+  if (auto s = Metadata::Decode(input); !s.ok()) {
+    return s;
+  }
+  if (!GetFixed32(input, &width)) {
+    return rocksdb::Status::InvalidArgument(kErrMetadataTooShort);
+  }
+  if (!GetFixed32(input, &depth)) {
+    return rocksdb::Status::InvalidArgument(kErrMetadataTooShort);
+  }
+  if (!GetFixed64(input, &counter)) {
+    return rocksdb::Status::InvalidArgument(kErrMetadataTooShort);
+  }
+
+  size_t array_size = width * depth;
+  array.resize(array_size);
+
+  for (size_t i = 0; i < array_size; ++i) {
+    uint32_t count = 0;
+    if (!GetFixed32(input, &count)) {
+      return rocksdb::Status::InvalidArgument(kErrMetadataTooShort);
+    }
+    array[i] = count;
+  }
+
+  return rocksdb::Status::OK();
+}

src/storage/redis_metadata.h

@@ -50,6 +50,7 @@ enum RedisType : uint8_t {
   kRedisBloomFilter = 9,
   kRedisJson = 10,
   kRedisHyperLogLog = 11,
+  kRedisCountMinSketch = 12,
 };
 
 struct RedisTypes {
@@ -91,9 +92,9 @@ enum RedisCommand {
   kRedisCmdLMove,
 };
 
-const std::vector<std::string> RedisTypeNames = {"none",   "string",    "hash",      "list",
-                                                 "set",    "zset",      "bitmap",    "sortedint",
-                                                 "stream", "MBbloom--", "ReJSON-RL", "hyperloglog"};
+const std::vector<std::string> RedisTypeNames = {"none",      "string",      "hash",          "list",   "set",
+                                                 "zset",      "bitmap",      "sortedint",     "stream", "MBbloom--",
+                                                 "ReJSON-RL", "hyperloglog", "countminsketch"};
 
 constexpr const char *kErrMsgWrongType = "WRONGTYPE Operation against a key holding the wrong kind of value";
 constexpr const char *kErrMsgKeyExpired = "the key was expired";
@@ -335,3 +336,15 @@ class HyperLogLogMetadata : public Metadata {
 
   EncodeType encode_type = EncodeType::DENSE;
 };
+
+class CountMinSketchMetadata : public Metadata {
+ public:
+  uint32_t width;
+  uint32_t depth;
+  uint64_t counter;
+  std::vector<uint32_t> array;
+
+  explicit CountMinSketchMetadata(bool generate_version = true) : Metadata(kRedisCountMinSketch, generate_version) {}
+  void Encode(std::string *dst) const override;
+  rocksdb::Status Decode(Slice *input) override;
+};