feat(cf): add CF.COUNT command

src/commands/cmd_cuckoo_filter.cc

@@ -131,8 +131,41 @@ class CommandCFAdd : public Commander {
   }
 };
 
-// Register the CF.RESERVE and CF.ADD commands
+class CommandCFCount : public Commander {
+ public:
+  Status Parse(const std::vector<std::string> &args) override {
+    // CF.COUNT key item
+    if (args.size() != 3) {
+      return {Status::RedisParseErr, errWrongNumOfArguments};
+    }
+    return Commander::Parse(args);
+  }
+
+  Status Execute(engine::Context &ctx, Server *srv, Connection *conn, std::string *output) override {
+    redis::CuckooChain cuckoo_db(srv->storage, conn->GetNamespace());
+    uint64_t count = 0;
+    auto s = cuckoo_db.Count(ctx, args_[1], args_[2], &count);
+
+    if (s.IsNotFound()) {
+      *output = redis::Integer(0);
+      return Status::OK();
+    }
+    if (!s.ok()) {
+      return {Status::RedisExecErr, s.ToString()};
+    }
+
+    *output = redis::Integer(count);
+    return Status::OK();
+  }
+};
+
+/* Register commands
+  CF.RESERVE
+  CF.ADD
+  CF.COUNT
+*/
 REDIS_REGISTER_COMMANDS(CuckooFilter, MakeCmdAttr<CommandCFReserve>("cf.reserve", -3, "write", 1, 1, 1),
-                        MakeCmdAttr<CommandCFAdd>("cf.add", 3, "write", 1, 1, 1))
+                        MakeCmdAttr<CommandCFAdd>("cf.add", 3, "write", 1, 1, 1),
+                        MakeCmdAttr<CommandCFCount>("cf.count", 3, "read-only", 1, 1, 1))
 
 }  // namespace redis

src/types/cuckoo_filter_page.cc

@@ -234,4 +234,20 @@ void CuckooPageCache::setBucketRefSlot(const BucketRef &bucket, uint32_t slot_id
   bucket.page->is_dirty = true;
 }
 
+rocksdb::Status CuckooPageCache::CountInBucket(uint16_t filter_index, uint32_t num_buckets, uint32_t bucket_index,
+                                               uint8_t fingerprint, uint32_t *count) {
+  *count = 0;
+  BucketRef bucket;
+  auto s = ensureBucketLoaded(filter_index, num_buckets, bucket_index, &bucket);
+  if (!s.ok()) return s;
+
+  for (size_t i = 0; i < bucket.size; ++i) {
+    size_t offset = bucket.offset + i;
+    if (static_cast<uint8_t>(bucket.page->data[offset]) == fingerprint) {
+      *count += 1;
+    }
+  }
+  return rocksdb::Status::OK();
+}
+
 }  // namespace redis

src/types/cuckoo_filter_page.h

@@ -49,6 +49,9 @@ class CuckooPageCache {
 
   void DiscardCachedPages();
 
+  rocksdb::Status CountInBucket(uint16_t filter_index, uint32_t num_buckets, uint32_t bucket_index, uint8_t fingerprint,
+                                uint32_t *count);
+
  private:
   struct PageEntry {
     std::string data;

src/types/cuckoo_filter_sub_filter.cc

@@ -97,6 +97,28 @@ rocksdb::Status CuckooSubFilter::WriteToBatch(rocksdb::WriteBatchBase *batch) {
   return pages_.WriteBackDirtyPages(batch);
 }
 
+rocksdb::Status CuckooSubFilter::Count(uint64_t hash, uint8_t fingerprint, uint32_t *count) {
+  uint32_t bucket1_idx = getPrimaryBucketIndex(hash);
+  uint32_t bucket2_idx = getSecondaryBucketIndex(hash, fingerprint);
+
+  auto s = pages_.PrefetchBuckets(filter_index_, num_buckets_, bucket1_idx, bucket2_idx);
+  if (!s.ok()) return s;
+
+  uint32_t c1 = 0, c2 = 0;
+  s = pages_.CountInBucket(filter_index_, num_buckets_, bucket1_idx, fingerprint, &c1);
+  if (!s.ok()) return s;
+
+  *count += c1;
+
+  if (bucket1_idx != bucket2_idx) {
+    s = pages_.CountInBucket(filter_index_, num_buckets_, bucket2_idx, fingerprint, &c2);
+    if (!s.ok()) return s;
+    *count += c2;
+  }
+
+  return rocksdb::Status::OK();
+}
+
 uint32_t CuckooSubFilter::getPrimaryBucketIndex(uint64_t hash) const { return hash % num_buckets_; }
 
 uint32_t CuckooSubFilter::getSecondaryBucketIndex(uint64_t hash, uint8_t fingerprint) const {

src/types/cuckoo_filter_sub_filter.h

@@ -44,6 +44,9 @@ class CuckooSubFilter {
   rocksdb::Status TryKickOutInsert(uint64_t hash, uint8_t fingerprint, uint16_t max_iterations, bool *inserted);
   rocksdb::Status WriteToBatch(rocksdb::WriteBatchBase *batch);
 
+  // Counts the occurrences of an item in this sub-filter by checking both primary and secondary buckets.
+  rocksdb::Status Count(uint64_t hash, uint8_t fingerprint, uint32_t *count);
+
  private:
   uint32_t getPrimaryBucketIndex(uint64_t hash) const;
   uint32_t getSecondaryBucketIndex(uint64_t hash, uint8_t fingerprint) const;

src/types/redis_cuckoo_chain.cc

@@ -179,6 +179,49 @@ rocksdb::Status CuckooChain::Add(engine::Context &ctx, const Slice &user_key, co
   return rocksdb::Status::Aborted("filter is full");
 }
 
+rocksdb::Status CuckooChain::Count(engine::Context &ctx, const Slice &user_key, const Slice &item, uint64_t *count) {
+  std::string ns_key = AppendNamespacePrefix(user_key);
+
+  CuckooChainMetadata metadata(false);
+  auto s = getCuckooChainMetadata(ctx, ns_key, &metadata);
+
+  if (s.IsNotFound()) {
+    return s;
+  }
+  if (!s.ok()) {
+    return s;
+  }
+
+  s = validateMetadata(metadata);
+  if (!s.ok()) {
+    return s;
+  }
+
+  uint64_t hash = CuckooFilterHelper::Hash(item.data(), item.size());
+  uint8_t fingerprint = CuckooFilterHelper::GenerateFingerprint(hash);
+
+  *count = 0;
+
+  for (uint16_t filter_idx = 0; filter_idx < metadata.n_filters; filter_idx++) {
+    uint32_t num_buckets = 0;
+    s = CuckooFilterHelper::GetFilterNumBuckets(metadata.base_capacity, metadata.expansion, metadata.bucket_size,
+                                                filter_idx, &num_buckets);
+    if (!s.ok()) {
+      return s;
+    }
+
+    CuckooSubFilter sub_filter(storage_, ctx, ns_key, storage_->IsSlotIdEncoded(), metadata.version,
+                               metadata.bucket_size, metadata.page_size, filter_idx, num_buckets);
+    uint32_t sub_count = 0;
+    s = sub_filter.Count(hash, fingerprint, &sub_count);
+    if (!s.ok()) {
+      return s;
+    }
+    *count += static_cast<uint64_t>(sub_count);
+  }
+  return rocksdb::Status::OK();
+}
+
 rocksdb::Status CuckooChain::tryCuckooInsert(engine::Context &ctx, const Slice &user_key, const std::string &ns_key,
                                              CuckooChainMetadata *metadata, uint64_t hash, uint8_t fingerprint,
                                              bool *inserted) {

src/types/redis_cuckoo_chain.h

@@ -47,6 +47,9 @@ class CuckooChain : public Database {
   // Duplicate items are allowed, so added is true whenever insertion succeeds.
   rocksdb::Status Add(engine::Context &ctx, const Slice &user_key, const Slice &item, bool *added);
 
+  // Returns the count of an item in the cuckoo filter, or 0 if not found.
+  rocksdb::Status Count(engine::Context &ctx, const Slice &user_key, const Slice &item, uint64_t *count);
+
  private:
   // Loads metadata for a cuckoo filter key.
   rocksdb::Status getCuckooChainMetadata(engine::Context &ctx, const Slice &ns_key, CuckooChainMetadata *metadata);

tests/gocase/unit/type/bloom/cuckoo_filter_test.go

@@ -187,4 +187,56 @@ func TestCuckooFilter(t *testing.T) {
 		require.NoError(t, result.Err())
 		require.Equal(t, int64(1), result.Val())
 	})
+
+	t.Run("Count no exists key returns zero", func(t *testing.T) {
+		key := "no_exist_key_count"
+		require.NoError(t, rdb.Del(ctx, key).Err())
+		require.Equal(t, int64(0), rdb.Do(ctx, "cf.count", key, "item").Val())
+	})
+
+	t.Run("Count empty filter", func(t *testing.T) {
+		key := "test_cf_count_empty"
+		require.NoError(t, rdb.Del(ctx, key).Err())
+		require.NoError(t, rdb.Do(ctx, "cf.reserve", key, "100").Err())
+		require.Equal(t, int64(0), rdb.Do(ctx, "cf.count", key, "item").Val())
+	})
+
+	t.Run("Count after add", func(t *testing.T) {
+		key := "test_cf_count_after_add"
+		require.NoError(t, rdb.Del(ctx, key).Err())
+		require.NoError(t, rdb.Do(ctx, "cf.reserve", key, "100").Err())
+		require.NoError(t, rdb.Do(ctx, "cf.add", key, "my_item").Err())
+		require.Equal(t, int64(1), rdb.Do(ctx, "cf.count", key, "my_item").Val())
+	})
+
+	t.Run("Count duplicate items", func(t *testing.T) {
+		key := "test_cf_count_dup"
+		require.NoError(t, rdb.Del(ctx, key).Err())
+		require.NoError(t, rdb.Do(ctx, "cf.reserve", key, "100").Err())
+		for i := 0; i < 5; i++ {
+			require.NoError(t, rdb.Do(ctx, "cf.add", key, "same_item").Err())
+		}
+		require.Equal(t, int64(5), rdb.Do(ctx, "cf.count", key, "same_item").Val())
+	})
+
+	t.Run("Count non-existent item", func(t *testing.T) {
+		key := "test_cf_count_not_exist_item"
+		require.NoError(t, rdb.Del(ctx, key).Err())
+		require.NoError(t, rdb.Do(ctx, "cf.reserve", key, "100").Err())
+		require.NoError(t, rdb.Do(ctx, "cf.add", key, "item1").Err())
+		require.Equal(t, int64(0), rdb.Do(ctx, "cf.count", key, "item2").Val())
+	})
+
+	t.Run("Count wrong type key", func(t *testing.T) {
+		key := "test_cf_count_wrong_type"
+		require.NoError(t, rdb.Del(ctx, key).Err())
+		require.NoError(t, rdb.Set(ctx, key, "value", 0).Err())
+		require.ErrorContains(t, rdb.Do(ctx, "cf.count", key, "item").Err(), "WRONGTYPE")
+	})
+
+	t.Run("Count wrong number of arguments", func(t *testing.T) {
+		require.Error(t, rdb.Do(ctx, "cf.count").Err())
+		require.Error(t, rdb.Do(ctx, "cf.count", "key_only").Err())
+		require.Error(t, rdb.Do(ctx, "cf.count", "key", "item1", "item2").Err())
+	})
 }