Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add the support of ZINTER and ZINTERCARD #1992

Merged
merged 6 commits into from
Jan 9, 2024
Merged

Add the support of ZINTER and ZINTERCARD #1992

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
c273151
Add the soupport of ZINTER and ZINTERCARD
kay011 Jan 7, 2024
5c9b2f8
delete unused include
kay011 Jan 7, 2024
990a105
Fix duplicated code
kay011 Jan 7, 2024
6733eb7
Fix duplicated code
kay011 Jan 7, 2024
019a790
Merge branch 'unstable' into feature_add_zinter
git-hulk Jan 8, 2024
abbfae5
Fix code style for lint
kay011 Jan 8, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
83 changes: 83 additions & 0 deletions src/commands/cmd_zset.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1252,6 +1252,87 @@ class CommandZInterStore : public CommandZUnionStore {
}
};

class CommandZInter : public CommandZUnion {
public:
CommandZInter() : CommandZUnion() {}

Status Execute(Server *srv, Connection *conn, std::string *output) override {
redis::ZSet zset_db(srv->storage, conn->GetNamespace());
std::vector<MemberScore> member_scores;
auto s = zset_db.Inter(keys_weights_, aggregate_method_, &member_scores);
if (!s.ok()) {
return {Status::RedisExecErr, s.ToString()};
}
auto ms_comparator = [](const MemberScore &ms1, const MemberScore &ms2) {
if (ms1.score == ms2.score) {
return ms1.member < ms2.member;
}
return ms1.score < ms2.score;
};
std::sort(member_scores.begin(), member_scores.end(), ms_comparator);
output->append(redis::MultiLen(member_scores.size() * (with_scores_ ? 2 : 1)));
for (const auto &member_score : member_scores) {
output->append(redis::BulkString(member_score.member));
if (with_scores_) output->append(redis::BulkString(util::Float2String(member_score.score)));
}
return Status::OK();
}

static CommandKeyRange Range(const std::vector<std::string> &args) {
int num_key = *ParseInt<int>(args[1], 10);
return {2, 1 + num_key, 1};
}
};

class CommandZInterCard : public Commander {
public:
Status Parse(const std::vector<std::string> &args) override {
CommandParser parser(args, 1);
numkeys_ = GET_OR_RET(parser.TakeInt<int>(NumericRange<int>{1, std::numeric_limits<int>::max()}));
for (size_t i = 0; i < numkeys_; ++i) {
keys_.emplace_back(GET_OR_RET(parser.TakeStr()));
}

// if set limit option
if (parser.Good()) {
if (parser.EatEqICase("limit")) {
auto res = parser.TakeInt<int64_t>();
if (!res.IsOK() || res.GetValue() < 0) {
return {Status::RedisParseErr, errLimitIsNegative};
}
limit_ = static_cast<size_t>(res.GetValue());
if (parser.Good()) {
return parser.InvalidSyntax();
}
} else {
return parser.InvalidSyntax();
}
}

return Commander::Parse(args);
}
Status Execute(Server *srv, Connection *conn, std::string *output) override {
redis::ZSet zset_db(srv->storage, conn->GetNamespace());
uint64_t count = 0;
auto s = zset_db.InterCard(keys_, limit_, &count);
if (!s.ok()) {
return {Status::RedisExecErr, s.ToString()};
}
*output = redis::Integer(count);
return Status::OK();
}

static CommandKeyRange Range(const std::vector<std::string> &args) {
int num_key = *ParseInt<int>(args[1], 10);
return {2, 1 + num_key, 1};
}

private:
size_t numkeys_ = 0;
size_t limit_ = 0;
std::vector<std::string> keys_;
};

class CommandZScan : public CommandSubkeyScanBase {
public:
CommandZScan() = default;
Expand Down Expand Up @@ -1281,6 +1362,8 @@ REDIS_REGISTER_COMMANDS(MakeCmdAttr<CommandZAdd>("zadd", -4, "write", 1, 1, 1),
MakeCmdAttr<CommandZCount>("zcount", 4, "read-only", 1, 1, 1),
MakeCmdAttr<CommandZIncrBy>("zincrby", 4, "write", 1, 1, 1),
MakeCmdAttr<CommandZInterStore>("zinterstore", -4, "write", CommandZInterStore::Range),
MakeCmdAttr<CommandZInter>("zinter", -3, "read-only", CommandZInter::Range),
MakeCmdAttr<CommandZInterCard>("zintercard", -3, "read-only", CommandZInterCard::Range),
MakeCmdAttr<CommandZLexCount>("zlexcount", 4, "read-only", 1, 1, 1),
MakeCmdAttr<CommandZPopMax>("zpopmax", -2, "write", 1, 1, 1),
MakeCmdAttr<CommandZPopMin>("zpopmin", -2, "write", 1, 1, 1),
Expand Down
46 changes: 46 additions & 0 deletions src/types/redis_zset.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -699,6 +699,52 @@ rocksdb::Status ZSet::Inter(const std::vector<KeyWeight> &keys_weights, Aggregat
return rocksdb::Status::OK();
}

rocksdb::Status ZSet::InterCard(const std::vector<std::string> &user_keys, uint64_t limit, uint64_t *inter_cnt) {
std::vector<std::string> lock_keys;
lock_keys.reserve(user_keys.size());
for (const auto &user_key : user_keys) {
std::string ns_key = AppendNamespacePrefix(user_key);
lock_keys.emplace_back(std::move(ns_key));
}
MultiLockGuard guard(storage_->GetLockManager(), lock_keys);

std::vector<MemberScores> mscores_list;
mscores_list.reserve(user_keys.size());
RangeScoreSpec spec;
for (const auto &user_key : user_keys) {
MemberScores mscores;
auto s = RangeByScore(user_key, spec, &mscores, nullptr);
if (!s.ok() || mscores.empty()) return s;
mscores_list.emplace_back(mscores);
}
std::sort(mscores_list.begin(), mscores_list.end(),
[](const MemberScores &v1, const MemberScores &v2) { return v1.size() < v2.size(); });

auto base_mscores = mscores_list[0];
std::map<std::string, size_t> member_counters;
uint64_t cardinality = 0;
for (const auto &base_ms : base_mscores) {
member_counters[base_ms.member] = 1;
for (size_t i = 1; i < mscores_list.size(); i++) {
for (const auto &ms : mscores_list[i]) {
if (base_ms.member == ms.member) {
member_counters[ms.member]++;
break;
}
}
}
if (member_counters[base_ms.member] == mscores_list.size()) {
cardinality++;
if (limit > 0 && cardinality >= limit) {
*inter_cnt = limit;
return rocksdb::Status::OK();
};
}
}
*inter_cnt = cardinality;
return rocksdb::Status::OK();
}

rocksdb::Status ZSet::UnionStore(const Slice &dst, const std::vector<KeyWeight> &keys_weights,
AggregateMethod aggregate_method, uint64_t *saved_cnt) {
*saved_cnt = 0;
Expand Down
1 change: 1 addition & 0 deletions src/types/redis_zset.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -111,6 +111,7 @@ class ZSet : public SubKeyScanner {
AggregateMethod aggregate_method, uint64_t *saved_cnt);
rocksdb::Status Inter(const std::vector<KeyWeight> &keys_weights, AggregateMethod aggregate_method,
std::vector<MemberScore> *members);
rocksdb::Status InterCard(const std::vector<std::string> &user_keys, uint64_t limit, uint64_t *inter_cnt);
rocksdb::Status UnionStore(const Slice &dst, const std::vector<KeyWeight> &keys_weights,
AggregateMethod aggregate_method, uint64_t *saved_cnt);
rocksdb::Status Union(const std::vector<KeyWeight> &keys_weights, AggregateMethod aggregate_method,
Expand Down
52 changes: 52 additions & 0 deletions tests/gocase/unit/type/zset/zset_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1236,6 +1236,58 @@ func basicTests(t *testing.T, rdb *redis.Client, ctx context.Context, encoding s
require.Equal(t, []redis.Z{{2, "b"}, {3, "c"}}, rdb.ZRangeWithScores(ctx, "zsetc", 0, -1).Val())
})

t.Run(fmt.Sprintf("ZINTER with AGGREGATE and WEIGHTS - %s", encoding), func(t *testing.T) {
createZset(rdb, ctx, "zseta", []redis.Z{
{Score: 1, Member: "a"},
{Score: 2, Member: "b"},
{Score: 3, Member: "c"},
})
createZset(rdb, ctx, "zsetb", []redis.Z{
{Score: 1, Member: "b"},
{Score: 2, Member: "c"},
{Score: 3, Member: "d"},
})

require.Equal(t, []string{"b", "c"}, rdb.ZInter(ctx, &redis.ZStore{Keys: []string{"zseta", "zsetb"}}).Val())
require.Equal(t, []string{"b", "c"}, rdb.ZInter(ctx, &redis.ZStore{Keys: []string{"zseta", "zsetb"}, Aggregate: "max"}).Val())
require.Equal(t, []string{"b", "c"}, rdb.ZInter(ctx, &redis.ZStore{Keys: []string{"zseta", "zsetb"}, Aggregate: "min"}).Val())
require.Equal(t, []string{"b", "c"}, rdb.ZInter(ctx, &redis.ZStore{Keys: []string{"zseta", "zsetb"}, Aggregate: "sum"}).Val())

require.Equal(t, []string{"b", "c"}, rdb.ZInter(ctx, &redis.ZStore{Keys: []string{"zseta", "zsetb"}, Weights: []float64{2, 3}, Aggregate: "sum"}).Val())
require.Equal(t, []string{"b", "c"}, rdb.ZInter(ctx, &redis.ZStore{Keys: []string{"zseta", "zsetb"}, Weights: []float64{2, 3}, Aggregate: "max"}).Val())
require.Equal(t, []string{"b", "c"}, rdb.ZInter(ctx, &redis.ZStore{Keys: []string{"zseta", "zsetb"}, Weights: []float64{2, 3}, Aggregate: "min"}).Val())

require.Equal(t, []redis.Z{{3, "b"}, {5, "c"}}, rdb.ZInterWithScores(ctx, &redis.ZStore{Keys: []string{"zseta", "zsetb"}}).Val())
require.Equal(t, []redis.Z{{2, "b"}, {3, "c"}}, rdb.ZInterWithScores(ctx, &redis.ZStore{Keys: []string{"zseta", "zsetb"}, Aggregate: "max"}).Val())
require.Equal(t, []redis.Z{{1, "b"}, {2, "c"}}, rdb.ZInterWithScores(ctx, &redis.ZStore{Keys: []string{"zseta", "zsetb"}, Aggregate: "min"}).Val())
require.Equal(t, []redis.Z{{3, "b"}, {5, "c"}}, rdb.ZInterWithScores(ctx, &redis.ZStore{Keys: []string{"zseta", "zsetb"}, Aggregate: "sum"}).Val())

require.Equal(t, []redis.Z{{7, "b"}, {12, "c"}}, rdb.ZInterWithScores(ctx, &redis.ZStore{Keys: []string{"zseta", "zsetb"}, Weights: []float64{2, 3}, Aggregate: "sum"}).Val())
require.Equal(t, []redis.Z{{4, "b"}, {6, "c"}}, rdb.ZInterWithScores(ctx, &redis.ZStore{Keys: []string{"zseta", "zsetb"}, Weights: []float64{2, 3}, Aggregate: "max"}).Val())
require.Equal(t, []redis.Z{{3, "b"}, {6, "c"}}, rdb.ZInterWithScores(ctx, &redis.ZStore{Keys: []string{"zseta", "zsetb"}, Weights: []float64{2, 3}, Aggregate: "min"}).Val())

require.Equal(t, 0, len(rdb.ZInter(ctx, &redis.ZStore{Keys: []string{"zseta", "zsetb", "zset_noexists"}}).Val()))
require.Equal(t, 0, len(rdb.ZInterWithScores(ctx, &redis.ZStore{Keys: []string{"zseta", "zsetb", "zset_noexists"}, Weights: []float64{2, 3}, Aggregate: "sum"}).Val()))
})

t.Run(fmt.Sprintf("ZINTERCARD with LIMIT - %s", encoding), func(t *testing.T) {
createZset(rdb, ctx, "zseta", []redis.Z{
{Score: 1, Member: "a"},
{Score: 2, Member: "b"},
{Score: 3, Member: "c"},
})
createZset(rdb, ctx, "zsetb", []redis.Z{
{Score: 1, Member: "b"},
{Score: 2, Member: "c"},
{Score: 3, Member: "d"},
})
require.Equal(t, int64(2), rdb.ZInterCard(ctx, 0, "zseta", "zsetb").Val())
require.Equal(t, int64(1), rdb.ZInterCard(ctx, 1, "zseta", "zsetb").Val())

require.Error(t, rdb.ZInterCard(ctx, -1, "zseta", "zsetb").Err())

})

for i, cmd := range []func(ctx context.Context, dest string, store *redis.ZStore) *redis.IntCmd{rdb.ZInterStore, rdb.ZUnionStore} {
var funcName string
switch i {
Expand Down