From def3e6364ba201d2b7160d175032960bbf770478 Mon Sep 17 00:00:00 2001 From: ChaseForward <3057706121@qq.com> Date: Sun, 27 Sep 2026 15:17:02 +0800 Subject: [PATCH] feat(cf): add CF.COUNT command Return the number of occurrences of an item in a cuckoo filter by scanning the primary and secondary buckets of every sub-filter. Returns 0 if the filter does not exist. Co-Authored-By: Claude Opus 5 (1M context) --- src/commands/cmd_cuckoo_filter.cc | 37 ++++++++++++- src/types/cuckoo_filter_page.cc | 16 ++++++ src/types/cuckoo_filter_page.h | 3 ++ src/types/cuckoo_filter_sub_filter.cc | 22 ++++++++ src/types/cuckoo_filter_sub_filter.h | 3 ++ src/types/redis_cuckoo_chain.cc | 43 +++++++++++++++ src/types/redis_cuckoo_chain.h | 3 ++ .../unit/type/bloom/cuckoo_filter_test.go | 52 +++++++++++++++++++ 8 files changed, 177 insertions(+), 2 deletions(-) diff --git a/src/commands/cmd_cuckoo_filter.cc b/src/commands/cmd_cuckoo_filter.cc index 2002912f94d..3950f1bf2b6 100644 --- a/src/commands/cmd_cuckoo_filter.cc +++ b/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 &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("cf.reserve", -3, "write", 1, 1, 1), - MakeCmdAttr("cf.add", 3, "write", 1, 1, 1)) + MakeCmdAttr("cf.add", 3, "write", 1, 1, 1), + MakeCmdAttr("cf.count", 3, "read-only", 1, 1, 1)) } // namespace redis diff --git a/src/types/cuckoo_filter_page.cc b/src/types/cuckoo_filter_page.cc index 1ef4a4c803d..3c55441d642 100644 --- a/src/types/cuckoo_filter_page.cc +++ b/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(bucket.page->data[offset]) == fingerprint) { + *count += 1; + } + } + return rocksdb::Status::OK(); +} + } // namespace redis diff --git a/src/types/cuckoo_filter_page.h b/src/types/cuckoo_filter_page.h index ff2efeb50c7..9c96ec7e0ca 100644 --- a/src/types/cuckoo_filter_page.h +++ b/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; diff --git a/src/types/cuckoo_filter_sub_filter.cc b/src/types/cuckoo_filter_sub_filter.cc index 0a0748cc354..08356c6948c 100644 --- a/src/types/cuckoo_filter_sub_filter.cc +++ b/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 { diff --git a/src/types/cuckoo_filter_sub_filter.h b/src/types/cuckoo_filter_sub_filter.h index 2bd26df8541..b778e143f70 100644 --- a/src/types/cuckoo_filter_sub_filter.h +++ b/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; diff --git a/src/types/redis_cuckoo_chain.cc b/src/types/redis_cuckoo_chain.cc index f7074ce9458..5b2f9f09dec 100644 --- a/src/types/redis_cuckoo_chain.cc +++ b/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(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) { diff --git a/src/types/redis_cuckoo_chain.h b/src/types/redis_cuckoo_chain.h index ae20b056cf3..e99b177e229 100644 --- a/src/types/redis_cuckoo_chain.h +++ b/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); diff --git a/tests/gocase/unit/type/bloom/cuckoo_filter_test.go b/tests/gocase/unit/type/bloom/cuckoo_filter_test.go index 852dde75ed0..dc630a77ab6 100644 --- a/tests/gocase/unit/type/bloom/cuckoo_filter_test.go +++ b/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()) + }) }