From 03b69e1c8fef7f0458236fb258d0df469294dbea Mon Sep 17 00:00:00 2001 From: donghao526 Date: Sun, 17 Aug 2025 22:50:25 +0800 Subject: [PATCH 01/43] feat: impl tdigest.revrank --- src/types/redis_tdigest.cc | 76 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 76 insertions(+) diff --git a/src/types/redis_tdigest.cc b/src/types/redis_tdigest.cc index f506ad92e2a..830c154567b 100644 --- a/src/types/redis_tdigest.cc +++ b/src/types/redis_tdigest.cc @@ -186,6 +186,82 @@ rocksdb::Status TDigest::Add(engine::Context& ctx, const Slice& digest_name, con return storage_->Write(ctx, storage_->DefaultWriteOptions(), batch->GetWriteBatch()); } +rocksdb::Status TDigest::RevRank(engine::Context& ctx, const Slice& digest_name, const std::vector& inputs, + std::vector* result) { + auto ns_key = AppendNamespacePrefix(digest_name); + TDigestMetadata metadata; + { + LockGuard guard(storage_->GetLockManager(), ns_key); + + if (auto status = getMetaDataByNsKey(ctx, ns_key, &metadata); !status.ok()) { + return status; + } + + if (metadata.total_observations == 0) { + result->resize(inputs.size(), 0); + return rocksdb::Status::OK(); + } + + if (metadata.unmerged_nodes > 0) { + auto batch = storage_->GetWriteBatchBase(); + WriteBatchLogData log_data(kRedisTDigest); + if (auto status = batch->PutLogData(log_data.Encode()); !status.ok()) { + return status; + } + + if (auto status = mergeCurrentBuffer(ctx, ns_key, batch, &metadata); !status.ok()) { + return status; + } + + std::string metadata_bytes; + metadata.Encode(&metadata_bytes); + if (auto status = batch->Put(metadata_cf_handle_, ns_key, metadata_bytes); !status.ok()) { + return status; + } + + if (auto status = storage_->Write(ctx, storage_->DefaultWriteOptions(), batch->GetWriteBatch()); !status.ok()) { + return status; + } + + ctx.RefreshLatestSnapshot(); + } + } + + std::vector centroids; + if (auto status = dumpCentroids(ctx, ns_key, metadata, ¢roids); !status.ok()) { + return status; + } + + auto dump_centroids = DummyCentroids(metadata, centroids); + + result->clear(); + result->reserve(inputs.size()); + + for (auto value : inputs) { + auto status_or_rank = TDigestRevRank(dump_centroids, value); + if (!status_or_rank) { + return rocksdb::Status::InvalidArgument(status_or_rank.Msg()); + } + result->push_back(*status_or_rank); + } + + return rocksdb::Status::OK(); +} + +StatusOr TDigestRevRank(const DummyCentroids& centroids, double value) { + double rank = 0; + auto it = centroids.Begin(); + while (it->Valid()) { + auto centroid_or = it->GetCentroid(); + if (!centroid_or) return {::Status::NotOK, centroid_or.Msg()}; + if (centroid_or->mean > value) { + rank += centroid_or->weight; + } + it->Next(); + } + return static_cast(rank); +} + rocksdb::Status TDigest::Quantile(engine::Context& ctx, const Slice& digest_name, const std::vector& qs, TDigestQuantitleResult* result) { auto ns_key = AppendNamespacePrefix(digest_name); From 97df4a1391b3e1f4f39b49aeba280d5d1d886782 Mon Sep 17 00:00:00 2001 From: donghao526 Date: Tue, 19 Aug 2025 19:24:30 +0800 Subject: [PATCH 02/43] feat: impl tdigest.revrank --- src/commands/cmd_tdigest.cc | 44 +++++++++++++++++++++++++++++++++++++ src/types/redis_tdigest.cc | 18 ++------------- src/types/redis_tdigest.h | 3 ++- src/types/tdigest.h | 19 ++++++++++++++++ 4 files changed, 67 insertions(+), 17 deletions(-) diff --git a/src/commands/cmd_tdigest.cc b/src/commands/cmd_tdigest.cc index e7815256d11..ca7e5398a81 100644 --- a/src/commands/cmd_tdigest.cc +++ b/src/commands/cmd_tdigest.cc @@ -176,6 +176,49 @@ class CommandTDigestAdd : public Commander { std::vector values_; }; +class CommandTDigestRevRank : public Commander { + public: + Status Parse(const std::vector &args) override { + key_name_ = args[1]; + inputs_.reserve(args.size() - 2); + for (size_t i = 2; i < args.size(); i++) { + auto value = ParseFloat(args[i]); + if (!value) { + return {Status::RedisParseErr, errValueIsNotFloat}; + } + inputs_.push_back(*value); + } + return Status::OK(); + } + Status Execute(engine::Context &ctx, Server *srv, Connection *conn, std::string *output) override { + TDigest tdigest(srv->storage, conn->GetNamespace()); + std::vector result; + result.reserve(inputs_.size()); + if (const auto s = tdigest.RevRank(ctx, key_name_, inputs_, &result); !s.ok()) { + if (s.IsNotFound()) { + return {Status::RedisExecErr, errKeyNotFound}; + } + return {Status::RedisExecErr, s.ToString()}; + } + + if (result.data()) { + std::vector rev_ranks; + rev_ranks.reserve(result.size()); + for (const auto v : result) { + rev_ranks.push_back(redis::Integer(v)); + } + *output = redis::Array(rev_ranks); + } else { + *output = redis::BulkString("nan"); + } + return Status::OK(); + } + + private: + std::string key_name_; + std::vector inputs_; +}; + class CommandTDigestMinMax : public Commander { public: explicit CommandTDigestMinMax(bool is_min) : is_min_(is_min) {} @@ -369,6 +412,7 @@ REDIS_REGISTER_COMMANDS(TDigest, MakeCmdAttr("tdigest.crea MakeCmdAttr("tdigest.add", -3, "write", 1, 1, 1), MakeCmdAttr("tdigest.max", 2, "read-only", 1, 1, 1), MakeCmdAttr("tdigest.min", 2, "read-only", 1, 1, 1), + MakeCmdAttr("tdigest.revrank", -3, "read-only", 1, 1, 1), MakeCmdAttr("tdigest.quantile", -3, "read-only", 1, 1, 1), MakeCmdAttr("tdigest.reset", 2, "write", 1, 1, 1), MakeCmdAttr("tdigest.merge", -4, "write", GetMergeKeyRange)); diff --git a/src/types/redis_tdigest.cc b/src/types/redis_tdigest.cc index 830c154567b..07727a9fbe4 100644 --- a/src/types/redis_tdigest.cc +++ b/src/types/redis_tdigest.cc @@ -187,7 +187,7 @@ rocksdb::Status TDigest::Add(engine::Context& ctx, const Slice& digest_name, con } rocksdb::Status TDigest::RevRank(engine::Context& ctx, const Slice& digest_name, const std::vector& inputs, - std::vector* result) { + std::vector* result) { auto ns_key = AppendNamespacePrefix(digest_name); TDigestMetadata metadata; { @@ -198,7 +198,7 @@ rocksdb::Status TDigest::RevRank(engine::Context& ctx, const Slice& digest_name, } if (metadata.total_observations == 0) { - result->resize(inputs.size(), 0); + result->resize(inputs.size(), -2); return rocksdb::Status::OK(); } @@ -244,23 +244,9 @@ rocksdb::Status TDigest::RevRank(engine::Context& ctx, const Slice& digest_name, } result->push_back(*status_or_rank); } - return rocksdb::Status::OK(); } -StatusOr TDigestRevRank(const DummyCentroids& centroids, double value) { - double rank = 0; - auto it = centroids.Begin(); - while (it->Valid()) { - auto centroid_or = it->GetCentroid(); - if (!centroid_or) return {::Status::NotOK, centroid_or.Msg()}; - if (centroid_or->mean > value) { - rank += centroid_or->weight; - } - it->Next(); - } - return static_cast(rank); -} rocksdb::Status TDigest::Quantile(engine::Context& ctx, const Slice& digest_name, const std::vector& qs, TDigestQuantitleResult* result) { diff --git a/src/types/redis_tdigest.h b/src/types/redis_tdigest.h index 2026cf94d05..3df670a4ca8 100644 --- a/src/types/redis_tdigest.h +++ b/src/types/redis_tdigest.h @@ -77,7 +77,8 @@ class TDigest : public SubKeyScanner { rocksdb::Status Merge(engine::Context& ctx, const Slice& dest_digest, const std::vector& source_digests, const TDigestMergeOptions& options); - + rocksdb::Status RevRank(engine::Context& ctx, const Slice& digest_name, const std::vector& inputs, + std::vector* result); rocksdb::Status GetMetaData(engine::Context& context, const Slice& digest_name, TDigestMetadata* metadata); private: diff --git a/src/types/tdigest.h b/src/types/tdigest.h index 1f416b48ecf..dcd3681ec7a 100644 --- a/src/types/tdigest.h +++ b/src/types/tdigest.h @@ -150,3 +150,22 @@ inline StatusOr TDigestQuantile(TD&& td, double q) { diff /= (lc.weight / 2 + rc.weight / 2); return Lerp(lc.mean, rc.mean, diff); } + +template +inline StatusOr TDigestRevRank(TD&& td, double value) { + if (value < td.Min()) { + return static_cast(td.TotalWeight()); + } + if (value > td.Max()) { + return static_cast(-1); + } + double rank = 0; + for (auto iter = td.Begin(); iter->Valid(); iter->Next()) { + if (auto centroid = GET_OR_RET(iter->GetCentroid()); centroid.mean > value) { + rank += centroid.weight; + } else if (centroid.mean == value) { + rank += centroid.weight / 2; + } + } + return static_cast(rank); +} \ No newline at end of file From 70a39d315314908d9949613ef6525e9374d756c5 Mon Sep 17 00:00:00 2001 From: donghao526 Date: Tue, 19 Aug 2025 19:24:30 +0800 Subject: [PATCH 03/43] feat: impl tdigest.revrank --- src/commands/cmd_tdigest.cc | 44 +++++++++++++++++++++++++++++++++++++ src/types/redis_tdigest.cc | 18 ++------------- src/types/redis_tdigest.h | 3 ++- src/types/tdigest.h | 19 ++++++++++++++++ 4 files changed, 67 insertions(+), 17 deletions(-) diff --git a/src/commands/cmd_tdigest.cc b/src/commands/cmd_tdigest.cc index e7815256d11..ca7e5398a81 100644 --- a/src/commands/cmd_tdigest.cc +++ b/src/commands/cmd_tdigest.cc @@ -176,6 +176,49 @@ class CommandTDigestAdd : public Commander { std::vector values_; }; +class CommandTDigestRevRank : public Commander { + public: + Status Parse(const std::vector &args) override { + key_name_ = args[1]; + inputs_.reserve(args.size() - 2); + for (size_t i = 2; i < args.size(); i++) { + auto value = ParseFloat(args[i]); + if (!value) { + return {Status::RedisParseErr, errValueIsNotFloat}; + } + inputs_.push_back(*value); + } + return Status::OK(); + } + Status Execute(engine::Context &ctx, Server *srv, Connection *conn, std::string *output) override { + TDigest tdigest(srv->storage, conn->GetNamespace()); + std::vector result; + result.reserve(inputs_.size()); + if (const auto s = tdigest.RevRank(ctx, key_name_, inputs_, &result); !s.ok()) { + if (s.IsNotFound()) { + return {Status::RedisExecErr, errKeyNotFound}; + } + return {Status::RedisExecErr, s.ToString()}; + } + + if (result.data()) { + std::vector rev_ranks; + rev_ranks.reserve(result.size()); + for (const auto v : result) { + rev_ranks.push_back(redis::Integer(v)); + } + *output = redis::Array(rev_ranks); + } else { + *output = redis::BulkString("nan"); + } + return Status::OK(); + } + + private: + std::string key_name_; + std::vector inputs_; +}; + class CommandTDigestMinMax : public Commander { public: explicit CommandTDigestMinMax(bool is_min) : is_min_(is_min) {} @@ -369,6 +412,7 @@ REDIS_REGISTER_COMMANDS(TDigest, MakeCmdAttr("tdigest.crea MakeCmdAttr("tdigest.add", -3, "write", 1, 1, 1), MakeCmdAttr("tdigest.max", 2, "read-only", 1, 1, 1), MakeCmdAttr("tdigest.min", 2, "read-only", 1, 1, 1), + MakeCmdAttr("tdigest.revrank", -3, "read-only", 1, 1, 1), MakeCmdAttr("tdigest.quantile", -3, "read-only", 1, 1, 1), MakeCmdAttr("tdigest.reset", 2, "write", 1, 1, 1), MakeCmdAttr("tdigest.merge", -4, "write", GetMergeKeyRange)); diff --git a/src/types/redis_tdigest.cc b/src/types/redis_tdigest.cc index 830c154567b..07727a9fbe4 100644 --- a/src/types/redis_tdigest.cc +++ b/src/types/redis_tdigest.cc @@ -187,7 +187,7 @@ rocksdb::Status TDigest::Add(engine::Context& ctx, const Slice& digest_name, con } rocksdb::Status TDigest::RevRank(engine::Context& ctx, const Slice& digest_name, const std::vector& inputs, - std::vector* result) { + std::vector* result) { auto ns_key = AppendNamespacePrefix(digest_name); TDigestMetadata metadata; { @@ -198,7 +198,7 @@ rocksdb::Status TDigest::RevRank(engine::Context& ctx, const Slice& digest_name, } if (metadata.total_observations == 0) { - result->resize(inputs.size(), 0); + result->resize(inputs.size(), -2); return rocksdb::Status::OK(); } @@ -244,23 +244,9 @@ rocksdb::Status TDigest::RevRank(engine::Context& ctx, const Slice& digest_name, } result->push_back(*status_or_rank); } - return rocksdb::Status::OK(); } -StatusOr TDigestRevRank(const DummyCentroids& centroids, double value) { - double rank = 0; - auto it = centroids.Begin(); - while (it->Valid()) { - auto centroid_or = it->GetCentroid(); - if (!centroid_or) return {::Status::NotOK, centroid_or.Msg()}; - if (centroid_or->mean > value) { - rank += centroid_or->weight; - } - it->Next(); - } - return static_cast(rank); -} rocksdb::Status TDigest::Quantile(engine::Context& ctx, const Slice& digest_name, const std::vector& qs, TDigestQuantitleResult* result) { diff --git a/src/types/redis_tdigest.h b/src/types/redis_tdigest.h index 2026cf94d05..3df670a4ca8 100644 --- a/src/types/redis_tdigest.h +++ b/src/types/redis_tdigest.h @@ -77,7 +77,8 @@ class TDigest : public SubKeyScanner { rocksdb::Status Merge(engine::Context& ctx, const Slice& dest_digest, const std::vector& source_digests, const TDigestMergeOptions& options); - + rocksdb::Status RevRank(engine::Context& ctx, const Slice& digest_name, const std::vector& inputs, + std::vector* result); rocksdb::Status GetMetaData(engine::Context& context, const Slice& digest_name, TDigestMetadata* metadata); private: diff --git a/src/types/tdigest.h b/src/types/tdigest.h index 1f416b48ecf..dcd3681ec7a 100644 --- a/src/types/tdigest.h +++ b/src/types/tdigest.h @@ -150,3 +150,22 @@ inline StatusOr TDigestQuantile(TD&& td, double q) { diff /= (lc.weight / 2 + rc.weight / 2); return Lerp(lc.mean, rc.mean, diff); } + +template +inline StatusOr TDigestRevRank(TD&& td, double value) { + if (value < td.Min()) { + return static_cast(td.TotalWeight()); + } + if (value > td.Max()) { + return static_cast(-1); + } + double rank = 0; + for (auto iter = td.Begin(); iter->Valid(); iter->Next()) { + if (auto centroid = GET_OR_RET(iter->GetCentroid()); centroid.mean > value) { + rank += centroid.weight; + } else if (centroid.mean == value) { + rank += centroid.weight / 2; + } + } + return static_cast(rank); +} \ No newline at end of file From 0d3e9cc8f5559a5acca4e7eef44efd44b8e8a551 Mon Sep 17 00:00:00 2001 From: donghao526 Date: Tue, 19 Aug 2025 23:46:18 +0800 Subject: [PATCH 04/43] feat: impl tdigest.revrank --- tests/cppunit/types/tdigest_test.cc | 77 +++++++++++++++++++++++++++++ 1 file changed, 77 insertions(+) diff --git a/tests/cppunit/types/tdigest_test.cc b/tests/cppunit/types/tdigest_test.cc index 91d1f311f4c..2ac75f1c70f 100644 --- a/tests/cppunit/types/tdigest_test.cc +++ b/tests/cppunit/types/tdigest_test.cc @@ -298,3 +298,80 @@ TEST_F(RedisTDigestTest, Quantile_returns_nan_on_empty_tdigest) { ASSERT_TRUE(status.ok()) << status.ToString(); ASSERT_FALSE(result.quantiles) << "should not have quantiles with empty tdigest"; } + +TEST_F(RedisTDigestTest, RevRank_on_the_set_contains_different_elements) { + std::string test_digest_name = "test_digest_revrank" + std::to_string(util::GetTimeStampMS()); + bool exists = false; + auto status = tdigest_->Create(*ctx_, test_digest_name, {100}, &exists); + ASSERT_FALSE(exists); + ASSERT_TRUE(status.ok()); + std::vector input{10, 20, 30, 40, 50, 60}; + status = tdigest_->Add(*ctx_, test_digest_name, input); + ASSERT_TRUE(status.ok()) << status.ToString(); + + std::vector result; + result.reserve(input.size()); + const std::vector value = {0, 10, 20, 30, 40, 50, 60, 70}; + status = tdigest_->RevRank(*ctx_, test_digest_name, value, &result); + const auto expect_result = std::vector{6, 5, 4, 3, 2, 1, 0, -1}; + + for (size_t i = 0; i < result.size(); i++) { + auto got = result[i]; + EXPECT_EQ(got, expect_result[i]); + } + ASSERT_TRUE(status.ok()) << status.ToString(); +} + +TEST_F(RedisTDigestTest, RevRank_on_the_set_contains_several_identical_elements) { + std::string test_digest_name = "test_digest_revrank" + std::to_string(util::GetTimeStampMS()); + bool exists = false; + auto status = tdigest_->Create(*ctx_, test_digest_name, {100}, &exists); + ASSERT_FALSE(exists); + ASSERT_TRUE(status.ok()); + std::vector input{10, 10, 10, 20, 20}; + status = tdigest_->Add(*ctx_, test_digest_name, input); + ASSERT_TRUE(status.ok()) << status.ToString(); + + std::vector result; + result.reserve(input.size()); + const std::vector value = {10, 20}; + status = tdigest_->RevRank(*ctx_, test_digest_name, value, &result); + const auto expect_result = std::vector{3, 1}; + for (size_t i = 0; i < result.size(); i++) { + auto got = result[i]; + EXPECT_EQ(got, expect_result[i]); + } + ASSERT_TRUE(status.ok()) << status.ToString(); + + status = tdigest_->Add(*ctx_, test_digest_name, std::vector{10}); + ASSERT_TRUE(status.ok()) << status.ToString(); + + result.clear(); + status = tdigest_->RevRank(*ctx_, test_digest_name, value, &result); + const auto expect_result_new = std::vector{4, 1}; + for (size_t i = 0; i < result.size(); i++) { + auto got = result[i]; + EXPECT_EQ(got, expect_result_new[i]); + } + ASSERT_TRUE(status.ok()) << status.ToString(); +} + +TEST_F(RedisTDigestTest, RevRank_on_empty_tdigest) { + std::string test_digest_name = "test_digest_revrank" + std::to_string(util::GetTimeStampMS()); + + bool exists = false; + auto status = tdigest_->Create(*ctx_, test_digest_name, {100}, &exists); + ASSERT_FALSE(exists); + ASSERT_TRUE(status.ok()); + + std::vector result; + result.reserve(2); + const std::vector value = {10, 20}; + status = tdigest_->RevRank(*ctx_, test_digest_name, value, &result); + const auto expect_result = std::vector{-2, -2}; + for (size_t i = 0; i < result.size(); i++) { + auto got = result[i]; + EXPECT_EQ(got, expect_result[i]); + } + ASSERT_TRUE(status.ok()) << status.ToString(); +} \ No newline at end of file From bb172a849256d8651499af5ee3ef0794540ecdcc Mon Sep 17 00:00:00 2001 From: donghao526 Date: Tue, 19 Aug 2025 23:46:18 +0800 Subject: [PATCH 05/43] test: add unit test for tdigest.revrank --- tests/cppunit/types/tdigest_test.cc | 76 +++++++++++++++++++++++++++++ 1 file changed, 76 insertions(+) diff --git a/tests/cppunit/types/tdigest_test.cc b/tests/cppunit/types/tdigest_test.cc index 91d1f311f4c..e543fb76e67 100644 --- a/tests/cppunit/types/tdigest_test.cc +++ b/tests/cppunit/types/tdigest_test.cc @@ -298,3 +298,79 @@ TEST_F(RedisTDigestTest, Quantile_returns_nan_on_empty_tdigest) { ASSERT_TRUE(status.ok()) << status.ToString(); ASSERT_FALSE(result.quantiles) << "should not have quantiles with empty tdigest"; } + +TEST_F(RedisTDigestTest, RevRank_on_the_set_contains_different_elements) { + std::string test_digest_name = "test_digest_revrank" + std::to_string(util::GetTimeStampMS()); + bool exists = false; + auto status = tdigest_->Create(*ctx_, test_digest_name, {100}, &exists); + ASSERT_FALSE(exists); + ASSERT_TRUE(status.ok()); + std::vector input{10, 20, 30, 40, 50, 60}; + status = tdigest_->Add(*ctx_, test_digest_name, input); + ASSERT_TRUE(status.ok()) << status.ToString(); + + std::vector result; + result.reserve(input.size()); + const std::vector value = {0, 10, 20, 30, 40, 50, 60, 70}; + status = tdigest_->RevRank(*ctx_, test_digest_name, value, &result); + const auto expect_result = std::vector{6, 5, 4, 3, 2, 1, 0, -1}; + + for (size_t i = 0; i < result.size(); i++) { + auto got = result[i]; + EXPECT_EQ(got, expect_result[i]); + } + ASSERT_TRUE(status.ok()) << status.ToString(); +} + +TEST_F(RedisTDigestTest, RevRank_on_the_set_contains_several_identical_elements) { + std::string test_digest_name = "test_digest_revrank" + std::to_string(util::GetTimeStampMS()); + bool exists = false; + auto status = tdigest_->Create(*ctx_, test_digest_name, {100}, &exists); + ASSERT_FALSE(exists); + ASSERT_TRUE(status.ok()); + std::vector input{10, 10, 10, 20, 20}; + status = tdigest_->Add(*ctx_, test_digest_name, input); + ASSERT_TRUE(status.ok()) << status.ToString(); + + std::vector result; + result.reserve(input.size()); + const std::vector value = {10, 20}; + status = tdigest_->RevRank(*ctx_, test_digest_name, value, &result); + const auto expect_result = std::vector{3, 1}; + for (size_t i = 0; i < result.size(); i++) { + auto got = result[i]; + EXPECT_EQ(got, expect_result[i]); + } + ASSERT_TRUE(status.ok()) << status.ToString(); + + status = tdigest_->Add(*ctx_, test_digest_name, std::vector{10}); + ASSERT_TRUE(status.ok()) << status.ToString(); + + result.clear(); + status = tdigest_->RevRank(*ctx_, test_digest_name, value, &result); + const auto expect_result_new = std::vector{4, 1}; + for (size_t i = 0; i < result.size(); i++) { + auto got = result[i]; + EXPECT_EQ(got, expect_result_new[i]); + } + ASSERT_TRUE(status.ok()) << status.ToString(); +} + +TEST_F(RedisTDigestTest, RevRank_on_empty_tdigest) { + std::string test_digest_name = "test_digest_revrank" + std::to_string(util::GetTimeStampMS()); + bool exists = false; + auto status = tdigest_->Create(*ctx_, test_digest_name, {100}, &exists); + ASSERT_FALSE(exists); + ASSERT_TRUE(status.ok()); + + std::vector result; + result.reserve(2); + const std::vector value = {10, 20}; + status = tdigest_->RevRank(*ctx_, test_digest_name, value, &result); + const auto expect_result = std::vector{-2, -2}; + for (size_t i = 0; i < result.size(); i++) { + auto got = result[i]; + EXPECT_EQ(got, expect_result[i]); + } + ASSERT_TRUE(status.ok()) << status.ToString(); +} \ No newline at end of file From 05d1202bff4e95eed8cec1fd41a924e8c0592925 Mon Sep 17 00:00:00 2001 From: donghao526 Date: Wed, 20 Aug 2025 08:39:09 +0800 Subject: [PATCH 06/43] add golang test cases for tdigest.revrank --- .../gocase/unit/type/tdigest/tdigest_test.go | 57 +++++++++++++++++++ 1 file changed, 57 insertions(+) diff --git a/tests/gocase/unit/type/tdigest/tdigest_test.go b/tests/gocase/unit/type/tdigest/tdigest_test.go index 3c2565ac066..63a127cdd9a 100644 --- a/tests/gocase/unit/type/tdigest/tdigest_test.go +++ b/tests/gocase/unit/type/tdigest/tdigest_test.go @@ -518,4 +518,61 @@ func tdigestTests(t *testing.T, configs util.KvrocksServerConfigs) { validation(newDestKey1) validation(newDestKey2) }) + + t.Run("tdigest.revrank with different arguments", func(t *testing.T) { + keyPrefix := "tdigest_revrank_" + + // Test invalid arguments + require.ErrorContains(t, rdb.Do(ctx, "TDIGEST.REVRANK").Err(), errMsgWrongNumberArg) + require.ErrorContains(t, rdb.Do(ctx, "TDIGEST.REVRANK", keyPrefix+"nonexistent").Err(), errMsgKeyNotExist) + + // Test with empty tdigest + key := keyPrefix + "test1" + require.NoError(t, rdb.Do(ctx, "TDIGEST.CREATE", key, "compression", "100").Err()) + rsp := rdb.Do(ctx, "TDIGEST.REVRANK", key) + require.NoError(t, rsp.Err()) + require.EqualValues(t, rsp.Val(), -2) + + // Test with set_contains several identical elements + require.NoError(t, rdb.Do(ctx, "TDIGEST.ADD", key, "10 10 10 20 20").Err()) + rsp = rdb.Do(ctx, "TDIGEST.REVRANK", key, "10 20") + require.NoError(t, rsp.Err()) + vals, err := rsp.Slice() + require.NoError(t, err) + require.Len(t, vals, 2) + expected := []int{3, 1} + for i, v := range vals { + rank, ok := v.(int) + require.True(t, ok, "expected string but got %T at index %d", v, i) + require.EqualValues(t, rank, expected[i]) + } + require.NoError(t, rdb.Do(ctx, "TDIGEST.ADD", key, "10").Err()) + rsp = rdb.Do(ctx, "TDIGEST.REVRANK", key, "10 20") + require.NoError(t, rsp.Err()) + vals, err = rsp.Slice() + require.NoError(t, err) + require.Len(t, vals, 2) + expected = []int{4, 1} + for i, v := range vals { + rank, ok := v.(int) + require.True(t, ok, "expected string but got %T at index %d", v, i) + require.EqualValues(t, rank, expected[i]) + } + + // Test with set_contains different elements + key2 := keyPrefix + "test2" + require.NoError(t, rdb.Do(ctx, "TDIGEST.CREATE", key2, "compression", "100").Err()) + require.NoError(t, rdb.Do(ctx, "TDIGEST.ADD", key2, "10 20 30 40 50 60").Err()) + rsp = rdb.Do(ctx, "TDIGEST.REVRANK", key2, "0 10 20 30 40 50 60 70") + require.NoError(t, rsp.Err()) + vals, err = rsp.Slice() + require.NoError(t, err) + require.Len(t, vals, 8) + expected = []int{6, 5, 4, 3, 2, 1, 0, -1} + for i, v := range vals { + rank, ok := v.(int) + require.True(t, ok, "expected string but got %T at index %d", v, i) + require.EqualValues(t, rank, expected[i]) + } + }) } From f688e1406e42c344789266a283bb855c3316b0e8 Mon Sep 17 00:00:00 2001 From: donghao526 Date: Wed, 20 Aug 2025 10:12:24 +0800 Subject: [PATCH 07/43] add golang test cases for tdigest.revrank --- tests/gocase/unit/type/tdigest/tdigest_test.go | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/gocase/unit/type/tdigest/tdigest_test.go b/tests/gocase/unit/type/tdigest/tdigest_test.go index 63a127cdd9a..f52a0fa6584 100644 --- a/tests/gocase/unit/type/tdigest/tdigest_test.go +++ b/tests/gocase/unit/type/tdigest/tdigest_test.go @@ -524,7 +524,10 @@ func tdigestTests(t *testing.T, configs util.KvrocksServerConfigs) { // Test invalid arguments require.ErrorContains(t, rdb.Do(ctx, "TDIGEST.REVRANK").Err(), errMsgWrongNumberArg) - require.ErrorContains(t, rdb.Do(ctx, "TDIGEST.REVRANK", keyPrefix+"nonexistent").Err(), errMsgKeyNotExist) + require.ErrorContains(t, rdb.Do(ctx, "TDIGEST.REVRANK", keyPrefix+"nonexistent").Err(), errMsgWrongNumberArg) + + // Test Non-existent key + require.ErrorContains(t, rdb.Do(ctx, "TDIGEST.REVRANK", keyPrefix+"nonexistent", "10").Err(), errMsgKeyNotExist) // Test with empty tdigest key := keyPrefix + "test1" From 8bcad0fb6a016f7ba80d578c4378cd2335a5b16a Mon Sep 17 00:00:00 2001 From: donghao526 Date: Wed, 20 Aug 2025 10:56:27 +0800 Subject: [PATCH 08/43] add golang test cases for tdigest.revrank --- tests/gocase/unit/type/tdigest/tdigest_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/gocase/unit/type/tdigest/tdigest_test.go b/tests/gocase/unit/type/tdigest/tdigest_test.go index f52a0fa6584..8362ff4174c 100644 --- a/tests/gocase/unit/type/tdigest/tdigest_test.go +++ b/tests/gocase/unit/type/tdigest/tdigest_test.go @@ -532,7 +532,7 @@ func tdigestTests(t *testing.T, configs util.KvrocksServerConfigs) { // Test with empty tdigest key := keyPrefix + "test1" require.NoError(t, rdb.Do(ctx, "TDIGEST.CREATE", key, "compression", "100").Err()) - rsp := rdb.Do(ctx, "TDIGEST.REVRANK", key) + rsp := rdb.Do(ctx, "TDIGEST.REVRANK", key, "10") require.NoError(t, rsp.Err()) require.EqualValues(t, rsp.Val(), -2) From 46ac984c6feff2b94d6c411775c37816fcfc8fe2 Mon Sep 17 00:00:00 2001 From: donghao526 Date: Wed, 20 Aug 2025 12:01:58 +0800 Subject: [PATCH 09/43] add golang test cases for tdigest.revrank --- .../gocase/unit/type/tdigest/tdigest_test.go | 30 ++++++++++++------- 1 file changed, 19 insertions(+), 11 deletions(-) diff --git a/tests/gocase/unit/type/tdigest/tdigest_test.go b/tests/gocase/unit/type/tdigest/tdigest_test.go index 8362ff4174c..3840a8bbe10 100644 --- a/tests/gocase/unit/type/tdigest/tdigest_test.go +++ b/tests/gocase/unit/type/tdigest/tdigest_test.go @@ -534,19 +534,27 @@ func tdigestTests(t *testing.T, configs util.KvrocksServerConfigs) { require.NoError(t, rdb.Do(ctx, "TDIGEST.CREATE", key, "compression", "100").Err()) rsp := rdb.Do(ctx, "TDIGEST.REVRANK", key, "10") require.NoError(t, rsp.Err()) - require.EqualValues(t, rsp.Val(), -2) + vals, err := rsp.Slice() + require.NoError(t, err) + require.Len(t, vals, 1) + expected := []int64{-2} + for i, v := range vals { + rank, ok := v.(int64) + require.True(t, ok, "expected int64 but got %T at index %d", v, i) + require.EqualValues(t, rank, expected[i]) + } // Test with set_contains several identical elements require.NoError(t, rdb.Do(ctx, "TDIGEST.ADD", key, "10 10 10 20 20").Err()) rsp = rdb.Do(ctx, "TDIGEST.REVRANK", key, "10 20") require.NoError(t, rsp.Err()) - vals, err := rsp.Slice() + vals, err = rsp.Slice() require.NoError(t, err) require.Len(t, vals, 2) - expected := []int{3, 1} + expected = []int64{3, 1} for i, v := range vals { - rank, ok := v.(int) - require.True(t, ok, "expected string but got %T at index %d", v, i) + rank, ok := v.(int64) + require.True(t, ok, "expected int64 but got %T at index %d", v, i) require.EqualValues(t, rank, expected[i]) } require.NoError(t, rdb.Do(ctx, "TDIGEST.ADD", key, "10").Err()) @@ -555,10 +563,10 @@ func tdigestTests(t *testing.T, configs util.KvrocksServerConfigs) { vals, err = rsp.Slice() require.NoError(t, err) require.Len(t, vals, 2) - expected = []int{4, 1} + expected = []int64{4, 1} for i, v := range vals { - rank, ok := v.(int) - require.True(t, ok, "expected string but got %T at index %d", v, i) + rank, ok := v.(int64) + require.True(t, ok, "expected int64 but got %T at index %d", v, i) require.EqualValues(t, rank, expected[i]) } @@ -571,10 +579,10 @@ func tdigestTests(t *testing.T, configs util.KvrocksServerConfigs) { vals, err = rsp.Slice() require.NoError(t, err) require.Len(t, vals, 8) - expected = []int{6, 5, 4, 3, 2, 1, 0, -1} + expected = []int64{6, 5, 4, 3, 2, 1, 0, -1} for i, v := range vals { - rank, ok := v.(int) - require.True(t, ok, "expected string but got %T at index %d", v, i) + rank, ok := v.(int64) + require.True(t, ok, "expected int64 but got %T at index %d", v, i) require.EqualValues(t, rank, expected[i]) } }) From 495e072044ab14a90a94552107b0fbc26b1a11be Mon Sep 17 00:00:00 2001 From: donghao526 Date: Wed, 20 Aug 2025 13:03:15 +0800 Subject: [PATCH 10/43] add golang test cases for tdigest.revrank --- tests/gocase/unit/type/tdigest/tdigest_test.go | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/tests/gocase/unit/type/tdigest/tdigest_test.go b/tests/gocase/unit/type/tdigest/tdigest_test.go index 3840a8bbe10..fd0d6c13b53 100644 --- a/tests/gocase/unit/type/tdigest/tdigest_test.go +++ b/tests/gocase/unit/type/tdigest/tdigest_test.go @@ -545,7 +545,7 @@ func tdigestTests(t *testing.T, configs util.KvrocksServerConfigs) { } // Test with set_contains several identical elements - require.NoError(t, rdb.Do(ctx, "TDIGEST.ADD", key, "10 10 10 20 20").Err()) + require.NoError(t, rdb.Do(ctx, "TDIGEST.ADD", key, 10, 10, 10, 20, 20).Err()) rsp = rdb.Do(ctx, "TDIGEST.REVRANK", key, "10 20") require.NoError(t, rsp.Err()) vals, err = rsp.Slice() @@ -557,8 +557,8 @@ func tdigestTests(t *testing.T, configs util.KvrocksServerConfigs) { require.True(t, ok, "expected int64 but got %T at index %d", v, i) require.EqualValues(t, rank, expected[i]) } - require.NoError(t, rdb.Do(ctx, "TDIGEST.ADD", key, "10").Err()) - rsp = rdb.Do(ctx, "TDIGEST.REVRANK", key, "10 20") + require.NoError(t, rdb.Do(ctx, "TDIGEST.ADD", key, 10).Err()) + rsp = rdb.Do(ctx, "TDIGEST.REVRANK", key, 10, 20) require.NoError(t, rsp.Err()) vals, err = rsp.Slice() require.NoError(t, err) @@ -573,8 +573,8 @@ func tdigestTests(t *testing.T, configs util.KvrocksServerConfigs) { // Test with set_contains different elements key2 := keyPrefix + "test2" require.NoError(t, rdb.Do(ctx, "TDIGEST.CREATE", key2, "compression", "100").Err()) - require.NoError(t, rdb.Do(ctx, "TDIGEST.ADD", key2, "10 20 30 40 50 60").Err()) - rsp = rdb.Do(ctx, "TDIGEST.REVRANK", key2, "0 10 20 30 40 50 60 70") + require.NoError(t, rdb.Do(ctx, "TDIGEST.ADD", key2, 10, 20, 30, 40, 50, 60).Err()) + rsp = rdb.Do(ctx, "TDIGEST.REVRANK", key2, 0, 10, 20, 30, 40, 50, 60, 70) require.NoError(t, rsp.Err()) vals, err = rsp.Slice() require.NoError(t, err) From 2b6785d55dccc95cdd2541ca339c8bff14059ba3 Mon Sep 17 00:00:00 2001 From: donghao526 Date: Wed, 20 Aug 2025 13:15:31 +0800 Subject: [PATCH 11/43] add golang test cases for tdigest.revrank --- tests/gocase/unit/type/tdigest/tdigest_test.go | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/tests/gocase/unit/type/tdigest/tdigest_test.go b/tests/gocase/unit/type/tdigest/tdigest_test.go index fd0d6c13b53..91c453bb33f 100644 --- a/tests/gocase/unit/type/tdigest/tdigest_test.go +++ b/tests/gocase/unit/type/tdigest/tdigest_test.go @@ -545,8 +545,8 @@ func tdigestTests(t *testing.T, configs util.KvrocksServerConfigs) { } // Test with set_contains several identical elements - require.NoError(t, rdb.Do(ctx, "TDIGEST.ADD", key, 10, 10, 10, 20, 20).Err()) - rsp = rdb.Do(ctx, "TDIGEST.REVRANK", key, "10 20") + require.NoError(t, rdb.Do(ctx, "TDIGEST.ADD", key, "10", "10", "10", "20", "20").Err()) + rsp = rdb.Do(ctx, "TDIGEST.REVRANK", key, "10", "20") require.NoError(t, rsp.Err()) vals, err = rsp.Slice() require.NoError(t, err) @@ -557,8 +557,9 @@ func tdigestTests(t *testing.T, configs util.KvrocksServerConfigs) { require.True(t, ok, "expected int64 but got %T at index %d", v, i) require.EqualValues(t, rank, expected[i]) } - require.NoError(t, rdb.Do(ctx, "TDIGEST.ADD", key, 10).Err()) - rsp = rdb.Do(ctx, "TDIGEST.REVRANK", key, 10, 20) + + require.NoError(t, rdb.Do(ctx, "TDIGEST.ADD", key, "10").Err()) + rsp = rdb.Do(ctx, "TDIGEST.REVRANK", key, "10", "20") require.NoError(t, rsp.Err()) vals, err = rsp.Slice() require.NoError(t, err) @@ -573,8 +574,8 @@ func tdigestTests(t *testing.T, configs util.KvrocksServerConfigs) { // Test with set_contains different elements key2 := keyPrefix + "test2" require.NoError(t, rdb.Do(ctx, "TDIGEST.CREATE", key2, "compression", "100").Err()) - require.NoError(t, rdb.Do(ctx, "TDIGEST.ADD", key2, 10, 20, 30, 40, 50, 60).Err()) - rsp = rdb.Do(ctx, "TDIGEST.REVRANK", key2, 0, 10, 20, 30, 40, 50, 60, 70) + require.NoError(t, rdb.Do(ctx, "TDIGEST.ADD", key2, "10", "20", "30", "40", "50", "60").Err()) + rsp = rdb.Do(ctx, "TDIGEST.REVRANK", key2, "0", "10", "20", "30", "40", "50", "60", "70") require.NoError(t, rsp.Err()) vals, err = rsp.Slice() require.NoError(t, err) From 3af3b54de679ba5fae40e18dd770320c4274a844 Mon Sep 17 00:00:00 2001 From: donghao526 Date: Wed, 20 Aug 2025 20:58:46 +0800 Subject: [PATCH 12/43] feat: impl tdigest.revrank --- src/types/redis_tdigest.cc | 77 ++++++++++++++++---------------------- src/types/redis_tdigest.h | 2 + 2 files changed, 35 insertions(+), 44 deletions(-) diff --git a/src/types/redis_tdigest.cc b/src/types/redis_tdigest.cc index f8015632eb6..4949b3f05e6 100644 --- a/src/types/redis_tdigest.cc +++ b/src/types/redis_tdigest.cc @@ -186,6 +186,35 @@ rocksdb::Status TDigest::Add(engine::Context& ctx, const Slice& digest_name, con return storage_->Write(ctx, storage_->DefaultWriteOptions(), batch->GetWriteBatch()); } +rocksdb::Status TDigest::mergeNodes(engine::Context& ctx, const std::string& ns_key, TDigestMetadata* metadata) { + if (metadata->unmerged_nodes == 0) { + return rocksdb::Status::OK(); + } + + auto batch = storage_->GetWriteBatchBase(); + WriteBatchLogData log_data(kRedisTDigest); + if (auto status = batch->PutLogData(log_data.Encode()); !status.ok()) { + return status; + } + + if (auto status = mergeCurrentBuffer(ctx, ns_key, batch, metadata); !status.ok()) { + return status; + } + + std::string metadata_bytes; + metadata->Encode(&metadata_bytes); + if (auto status = batch->Put(metadata_cf_handle_, ns_key, metadata_bytes); !status.ok()) { + return status; + } + + if (auto status = storage_->Write(ctx, storage_->DefaultWriteOptions(), batch->GetWriteBatch()); !status.ok()) { + return status; + } + + ctx.RefreshLatestSnapshot(); + return rocksdb::Status::OK(); +} + rocksdb::Status TDigest::RevRank(engine::Context& ctx, const Slice& digest_name, const std::vector& inputs, std::vector* result) { auto ns_key = AppendNamespacePrefix(digest_name); @@ -202,28 +231,8 @@ rocksdb::Status TDigest::RevRank(engine::Context& ctx, const Slice& digest_name, return rocksdb::Status::OK(); } - if (metadata.unmerged_nodes > 0) { - auto batch = storage_->GetWriteBatchBase(); - WriteBatchLogData log_data(kRedisTDigest); - if (auto status = batch->PutLogData(log_data.Encode()); !status.ok()) { - return status; - } - - if (auto status = mergeCurrentBuffer(ctx, ns_key, batch, &metadata); !status.ok()) { - return status; - } - - std::string metadata_bytes; - metadata.Encode(&metadata_bytes); - if (auto status = batch->Put(metadata_cf_handle_, ns_key, metadata_bytes); !status.ok()) { - return status; - } - - if (auto status = storage_->Write(ctx, storage_->DefaultWriteOptions(), batch->GetWriteBatch()); !status.ok()) { - return status; - } - - ctx.RefreshLatestSnapshot(); + if (auto status = mergeNodes(ctx, ns_key, &metadata); !status.ok()) { + return status; } } @@ -262,28 +271,8 @@ rocksdb::Status TDigest::Quantile(engine::Context& ctx, const Slice& digest_name return rocksdb::Status::OK(); } - if (metadata.unmerged_nodes > 0) { - auto batch = storage_->GetWriteBatchBase(); - WriteBatchLogData log_data(kRedisTDigest); - if (auto status = batch->PutLogData(log_data.Encode()); !status.ok()) { - return status; - } - - if (auto status = mergeCurrentBuffer(ctx, ns_key, batch, &metadata); !status.ok()) { - return status; - } - - std::string metadata_bytes; - metadata.Encode(&metadata_bytes); - if (auto status = batch->Put(metadata_cf_handle_, ns_key, metadata_bytes); !status.ok()) { - return status; - } - - if (auto status = storage_->Write(ctx, storage_->DefaultWriteOptions(), batch->GetWriteBatch()); !status.ok()) { - return status; - } - - ctx.RefreshLatestSnapshot(); + if (auto status = mergeNodes(ctx, ns_key, &metadata); !status.ok()) { + return status; } } diff --git a/src/types/redis_tdigest.h b/src/types/redis_tdigest.h index 3df670a4ca8..beb3d6933b6 100644 --- a/src/types/redis_tdigest.h +++ b/src/types/redis_tdigest.h @@ -118,6 +118,8 @@ class TDigest : public SubKeyScanner { std::string internalSegmentGuardPrefixKey(const TDigestMetadata& metadata, const std::string& ns_key, SegmentType seg) const; + rocksdb::Status mergeNodes(engine::Context& ctx, const std::string& ns_key, TDigestMetadata* metadata); + rocksdb::Status mergeCurrentBuffer(engine::Context& ctx, const std::string& ns_key, ObserverOrUniquePtr& batch, TDigestMetadata* metadata, const std::vector* additional_buffer = nullptr, From eb8674f5c370d53a17f39a08f3ff1c22c37a9fa8 Mon Sep 17 00:00:00 2001 From: donghao526 Date: Thu, 28 Aug 2025 00:17:59 +0800 Subject: [PATCH 13/43] feat: impl tdigest.revrank --- src/commands/cmd_tdigest.cc | 2 +- src/types/redis_tdigest.cc | 19 +++----- src/types/redis_tdigest.h | 2 +- src/types/tdigest.h | 74 ++++++++++++++++++++++++----- tests/cppunit/types/tdigest_test.cc | 8 ++-- 5 files changed, 74 insertions(+), 31 deletions(-) diff --git a/src/commands/cmd_tdigest.cc b/src/commands/cmd_tdigest.cc index ca7e5398a81..4bc1f2df9c8 100644 --- a/src/commands/cmd_tdigest.cc +++ b/src/commands/cmd_tdigest.cc @@ -194,7 +194,7 @@ class CommandTDigestRevRank : public Commander { TDigest tdigest(srv->storage, conn->GetNamespace()); std::vector result; result.reserve(inputs_.size()); - if (const auto s = tdigest.RevRank(ctx, key_name_, inputs_, &result); !s.ok()) { + if (const auto s = tdigest.RevRank(ctx, key_name_, inputs_, result); !s.ok()) { if (s.IsNotFound()) { return {Status::RedisExecErr, errKeyNotFound}; } diff --git a/src/types/redis_tdigest.cc b/src/types/redis_tdigest.cc index 4949b3f05e6..feded24fcae 100644 --- a/src/types/redis_tdigest.cc +++ b/src/types/redis_tdigest.cc @@ -79,6 +79,8 @@ class DummyCentroids { return Valid(); } bool Valid() const { return iter_ != centroids_.cend(); } + bool IsAtBegin() const { return iter_ == centroids_.cbegin(); } + StatusOr GetCentroid() const { if (iter_ == centroids_.cend()) { return {::Status::NotOK, "invalid iterator during decoding tdigest centroid"}; @@ -216,7 +218,7 @@ rocksdb::Status TDigest::mergeNodes(engine::Context& ctx, const std::string& ns_ } rocksdb::Status TDigest::RevRank(engine::Context& ctx, const Slice& digest_name, const std::vector& inputs, - std::vector* result) { + std::vector& result) { auto ns_key = AppendNamespacePrefix(digest_name); TDigestMetadata metadata; { @@ -227,7 +229,7 @@ rocksdb::Status TDigest::RevRank(engine::Context& ctx, const Slice& digest_name, } if (metadata.total_observations == 0) { - result->resize(inputs.size(), -2); + result.resize(inputs.size(), -2); return rocksdb::Status::OK(); } @@ -242,16 +244,9 @@ rocksdb::Status TDigest::RevRank(engine::Context& ctx, const Slice& digest_name, } auto dump_centroids = DummyCentroids(metadata, centroids); - - result->clear(); - result->reserve(inputs.size()); - - for (auto value : inputs) { - auto status_or_rank = TDigestRevRank(dump_centroids, value); - if (!status_or_rank) { - return rocksdb::Status::InvalidArgument(status_or_rank.Msg()); - } - result->push_back(*status_or_rank); + auto status = TDigestRank(dump_centroids, inputs, result); + if (!status) { + return rocksdb::Status::InvalidArgument(status.Msg()); } return rocksdb::Status::OK(); } diff --git a/src/types/redis_tdigest.h b/src/types/redis_tdigest.h index beb3d6933b6..02ecc24c016 100644 --- a/src/types/redis_tdigest.h +++ b/src/types/redis_tdigest.h @@ -78,7 +78,7 @@ class TDigest : public SubKeyScanner { rocksdb::Status Merge(engine::Context& ctx, const Slice& dest_digest, const std::vector& source_digests, const TDigestMergeOptions& options); rocksdb::Status RevRank(engine::Context& ctx, const Slice& digest_name, const std::vector& inputs, - std::vector* result); + std::vector& result); rocksdb::Status GetMetaData(engine::Context& context, const Slice& digest_name, TDigestMetadata* metadata); private: diff --git a/src/types/tdigest.h b/src/types/tdigest.h index dcd3681ec7a..b096187ba29 100644 --- a/src/types/tdigest.h +++ b/src/types/tdigest.h @@ -22,6 +22,7 @@ #include +#include #include #include "common/status.h" @@ -80,6 +81,9 @@ class TDSample { // reference: // https://github.com/apache/arrow/blob/27bbd593625122a4a25d9471c8aaf5df54a6dcf9/cpp/src/arrow/util/tdigest.cc#L38 static inline double Lerp(double a, double b, double t) { return a + t * (b - a); } +// static inline int CalculateRank(int total_weight, double cumulative_weight, bool reverse) { +// return reverse ? total_weight - 1 - static_cast(cumulative_weight) : static_cast(cumulative_weight); +// } template inline StatusOr TDigestQuantile(TD&& td, double q) { @@ -152,20 +156,64 @@ inline StatusOr TDigestQuantile(TD&& td, double q) { } template -inline StatusOr TDigestRevRank(TD&& td, double value) { - if (value < td.Min()) { - return static_cast(td.TotalWeight()); - } - if (value > td.Max()) { - return static_cast(-1); +inline Status TDigestRank(TD&& td, const std::vector& inputs, std::vector& result) { + std::vector indices(inputs.size()); + std::iota(indices.begin(), indices.end(), 0); + std::sort(indices.begin(), indices.end(), [&inputs](size_t a, size_t b) { return inputs[a] < inputs[b]; }); + + result.resize(inputs.size()); + + size_t i = indices.size(); + double cumulative_weight = 0; + + // handle inputs larger than maximum + while (i > 0 && inputs[indices[i - 1]] > td.Max()) { + result[indices[i - 1]] = -1; + i--; } - double rank = 0; - for (auto iter = td.Begin(); iter->Valid(); iter->Next()) { - if (auto centroid = GET_OR_RET(iter->GetCentroid()); centroid.mean > value) { - rank += centroid.weight; - } else if (centroid.mean == value) { - rank += centroid.weight / 2; + + // reverse iterate through centroids and calculate reverse rank for each input + auto iter = td.End(); + while (i > 0) { + auto centroid = GET_OR_RET(iter->GetCentroid()); + if (centroid.mean > inputs[indices[i - 1]]) { + cumulative_weight += centroid.weight; + } else if (centroid.mean == inputs[indices[i - 1]]) { + auto current_mean_cumulative_weight = cumulative_weight + centroid.weight / 2; + cumulative_weight += centroid.weight; + auto current_mean = centroid.mean; + // cumulative all the centroids which has the same mean + while (!iter->IsAtBegin() && iter->Prev()) { + auto next_centroid = GET_OR_RET(iter->GetCentroid()); + if (current_mean != next_centroid.mean) { + // move back to the last equal centroid, because we will process it in the next loop + iter->Next(); + break; + } + current_mean_cumulative_weight += centroid.weight / 2; + cumulative_weight += centroid.weight; + } + + result[indices[i - 1]] = static_cast(current_mean_cumulative_weight); + i--; + while ((i > 0) && (inputs[indices[i]] == inputs[indices[i - 1]])) { + result[indices[i - 1]] = result[indices[i]]; + i--; + } + } else { + result[indices[i - 1]] = static_cast(cumulative_weight); + i--; } + if (iter->IsAtBegin()) { + break; + } + iter->Prev(); + } + + // handle inputs less than minimum + while (i > 0) { + result[indices[i - 1]] = static_cast(td.TotalWeight()); + i--; } - return static_cast(rank); + return Status::OK(); } \ No newline at end of file diff --git a/tests/cppunit/types/tdigest_test.cc b/tests/cppunit/types/tdigest_test.cc index e543fb76e67..09b27fd70cd 100644 --- a/tests/cppunit/types/tdigest_test.cc +++ b/tests/cppunit/types/tdigest_test.cc @@ -312,7 +312,7 @@ TEST_F(RedisTDigestTest, RevRank_on_the_set_contains_different_elements) { std::vector result; result.reserve(input.size()); const std::vector value = {0, 10, 20, 30, 40, 50, 60, 70}; - status = tdigest_->RevRank(*ctx_, test_digest_name, value, &result); + status = tdigest_->RevRank(*ctx_, test_digest_name, value, result); const auto expect_result = std::vector{6, 5, 4, 3, 2, 1, 0, -1}; for (size_t i = 0; i < result.size(); i++) { @@ -335,7 +335,7 @@ TEST_F(RedisTDigestTest, RevRank_on_the_set_contains_several_identical_elements) std::vector result; result.reserve(input.size()); const std::vector value = {10, 20}; - status = tdigest_->RevRank(*ctx_, test_digest_name, value, &result); + status = tdigest_->RevRank(*ctx_, test_digest_name, value, result); const auto expect_result = std::vector{3, 1}; for (size_t i = 0; i < result.size(); i++) { auto got = result[i]; @@ -347,7 +347,7 @@ TEST_F(RedisTDigestTest, RevRank_on_the_set_contains_several_identical_elements) ASSERT_TRUE(status.ok()) << status.ToString(); result.clear(); - status = tdigest_->RevRank(*ctx_, test_digest_name, value, &result); + status = tdigest_->RevRank(*ctx_, test_digest_name, value, result); const auto expect_result_new = std::vector{4, 1}; for (size_t i = 0; i < result.size(); i++) { auto got = result[i]; @@ -366,7 +366,7 @@ TEST_F(RedisTDigestTest, RevRank_on_empty_tdigest) { std::vector result; result.reserve(2); const std::vector value = {10, 20}; - status = tdigest_->RevRank(*ctx_, test_digest_name, value, &result); + status = tdigest_->RevRank(*ctx_, test_digest_name, value, result); const auto expect_result = std::vector{-2, -2}; for (size_t i = 0; i < result.size(); i++) { auto got = result[i]; From 543fda09ee6d888ea3fd4a758ed2fd821ed974a1 Mon Sep 17 00:00:00 2001 From: Zhixin Wen Date: Tue, 19 Aug 2025 18:31:46 -0700 Subject: [PATCH 14/43] fix(replication): Fix Seg Fault On Signal When Replication is Enabled (#3131) --- src/common/thread_util.h | 10 ++++++++++ utils/kvrocks2redis/redis_writer.cc | 17 +++++++---------- 2 files changed, 17 insertions(+), 10 deletions(-) diff --git a/src/common/thread_util.h b/src/common/thread_util.h index 67b1d0614b9..f4c9cc75393 100644 --- a/src/common/thread_util.h +++ b/src/common/thread_util.h @@ -20,6 +20,8 @@ #pragma once +#include + #include #include @@ -34,6 +36,14 @@ template StatusOr CreateThread(const char *name, F f) { try { return std::thread([name, f = std::move(f)] { + // Block SIGTERM and SIGINT signals in this thread, + // the signal should be handled by the main thread or a dedicated thread. + sigset_t signal_set; + sigemptyset(&signal_set); + sigaddset(&signal_set, SIGTERM); + sigaddset(&signal_set, SIGINT); + pthread_sigmask(SIG_BLOCK, &signal_set, nullptr); + ThreadSetName(name); f(); }); diff --git a/utils/kvrocks2redis/redis_writer.cc b/utils/kvrocks2redis/redis_writer.cc index 932d96ca98a..98f75e43250 100644 --- a/utils/kvrocks2redis/redis_writer.cc +++ b/utils/kvrocks2redis/redis_writer.cc @@ -24,23 +24,20 @@ #include #include -#include - #include "io_util.h" #include "server/redis_reply.h" #include "thread_util.h" RedisWriter::RedisWriter(kvrocks2redis::Config *config) : Writer(config) { - try { - t_ = std::thread([this]() { - util::ThreadSetName("redis-writer"); - this->sync(); - assert(stop_flag_); - }); - } catch (const std::system_error &e) { - error("Failed to create thread: {}", e.what()); + auto t = util::CreateThread("redis-writer", [this]() { + this->sync(); + assert(stop_flag_); + }); + if (!t.IsOK()) { + error("Failed to create thread: {}", t.Msg()); return; } + t_ = std::move(*t); } RedisWriter::~RedisWriter() { From e0d39a7c39cb8305f384282bc6fb207a809676ee Mon Sep 17 00:00:00 2001 From: Twice Date: Wed, 20 Aug 2025 09:59:45 +0800 Subject: [PATCH 15/43] chore(.asf.yaml): make 2.13 a protected branches (#3129) Co-authored-by: hulk --- .asf.yaml | 1 + 1 file changed, 1 insertion(+) diff --git a/.asf.yaml b/.asf.yaml index 4f8b03c46b0..9034a87aa1c 100644 --- a/.asf.yaml +++ b/.asf.yaml @@ -56,6 +56,7 @@ github: '2.10': {} '2.11': {} '2.12': {} + '2.13': {} notifications: commits: commits@kvrocks.apache.org From a4ed14c9ae0c124508c8ebff06ba837a14c15ed7 Mon Sep 17 00:00:00 2001 From: RX Xiao <74295100+yezhizi@users.noreply.github.com> Date: Wed, 20 Aug 2025 12:57:55 +0800 Subject: [PATCH 16/43] feat(ts): Add support for data writing and `TS.CREATE`, `TS.ADD/MADD` commands (#3107) Co-authored-by: Twice --- src/commands/cmd_timeseries.cc | 329 ++++++++++++++++++ src/commands/commander.h | 1 + src/types/redis_timeseries.cc | 246 ++++++++++++- src/types/redis_timeseries.h | 47 ++- src/types/timeseries.cc | 122 +++++-- src/types/timeseries.h | 30 +- tests/cppunit/types/timeseries_chunk_test.cc | 133 ++++++- tests/cppunit/types/timeseries_test.cc | 106 ++++++ .../unit/type/timeseries/timeseries_test.go | 147 ++++++++ 9 files changed, 1108 insertions(+), 53 deletions(-) create mode 100644 src/commands/cmd_timeseries.cc create mode 100644 tests/cppunit/types/timeseries_test.cc create mode 100644 tests/gocase/unit/type/timeseries/timeseries_test.go diff --git a/src/commands/cmd_timeseries.cc b/src/commands/cmd_timeseries.cc new file mode 100644 index 00000000000..b993d3b0c5a --- /dev/null +++ b/src/commands/cmd_timeseries.cc @@ -0,0 +1,329 @@ +/* + * 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 "command_parser.h" +#include "commander.h" +#include "error_constants.h" +#include "server/server.h" +#include "types/redis_timeseries.h" + +namespace { +constexpr const char *errBadRetention = "Couldn't parse RETENTION"; +constexpr const char *errBadChunkSize = "invalid CHUNK_SIZE"; +constexpr const char *errBadEncoding = "unknown ENCODING parameter"; +constexpr const char *errDuplicatePolicy = "Unknown DUPLICATE_POLICY"; +constexpr const char *errInvalidTimestamp = "invalid timestamp"; +constexpr const char *errInvalidValue = "invalid value"; +constexpr const char *errOldTimestamp = "Timestamp is older than retention"; +constexpr const char *errDupBlock = + "Error at upsert, update is not supported when DUPLICATE_POLICY is set to BLOCK mode"; +constexpr const char *errTSKeyNotFound = "the key is not a TSDB key"; + +std::string FormatAddResultAsRedisReply(TSChunk::AddResultWithTS res) { + using AddResult = TSChunk::AddResult; + switch (res.first) { + case AddResult::kOk: + return redis::Integer(res.second); + case AddResult::kOld: + return redis::Error({Status::NotOK, errOldTimestamp}); + case AddResult::kBlock: + return redis::Error({Status::NotOK, errDupBlock}); + default: + unreachable(); + } + return ""; +} + +} // namespace + +namespace redis { + +class KeywordCommandBase : public Commander { + public: + KeywordCommandBase(size_t skip_num, size_t tail_skip_num) : skip_num_(skip_num), tail_skip_num_(tail_skip_num) {} + + Status Parse(const std::vector &args) override { + TSOptionsParser parser(std::next(args.begin(), static_cast(skip_num_)), + std::prev(args.end(), static_cast(tail_skip_num_))); + + while (parser.Good()) { + bool handled = false; + for (const auto &handler : handlers_) { + if (parser.EatEqICase(handler.first)) { + Status s = handler.second(parser); + if (!s.IsOK()) return s; + handled = true; + break; + } + } + + if (!handled) { + parser.Skip(1); + } + } + + return Commander::Parse(args); + } + + protected: + using TSOptionsParser = CommandParser; + + template + void registerHandler(const std::string &keyword, Handler &&handler) { + handlers_.emplace_back(keyword, std::forward(handler)); + } + + void setSkipNum(size_t num) { skip_num_ = num; } + + void setTailSkipNum(size_t num) { tail_skip_num_ = num; } + + private: + size_t skip_num_ = 0; + size_t tail_skip_num_ = 0; + + std::vector>> handlers_; +}; + +class CommandTSCreateBase : public KeywordCommandBase { + public: + CommandTSCreateBase(size_t skip_num, size_t tail_skip_num) : KeywordCommandBase(skip_num, tail_skip_num) { + registerHandler("RETENTION", [this](TSOptionsParser &parser) { return handleRetention(parser); }); + registerHandler("CHUNK_SIZE", [this](TSOptionsParser &parser) { return handleChunkSize(parser); }); + registerHandler("ENCODING", [this](TSOptionsParser &parser) { return handleEncoding(parser); }); + registerHandler("DUPLICATE_POLICY", [this](TSOptionsParser &parser) { return handleDuplicatePolicy(parser); }); + registerHandler("LABELS", [this](TSOptionsParser &parser) { return handleLabels(parser); }); + } + + protected: + using DuplicatePolicy = TimeSeriesMetadata::DuplicatePolicy; + + const TSCreateOption &getCreateOption() const { return create_option_; } + + private: + TSCreateOption create_option_; + + Status handleRetention(TSOptionsParser &parser) { + auto parse_retention = parser.TakeInt(); + if (!parse_retention.IsOK()) { + return {Status::RedisParseErr, errBadRetention}; + } + create_option_.retention_time = parse_retention.GetValue(); + return Status::OK(); + } + + Status handleChunkSize(TSOptionsParser &parser) { + auto parse_chunk_size = parser.TakeInt(); + if (!parse_chunk_size.IsOK()) { + return {Status::RedisParseErr, errBadChunkSize}; + } + create_option_.chunk_size = parse_chunk_size.GetValue(); + return Status::OK(); + } + + Status handleEncoding(TSOptionsParser &parser) { + using ChunkType = TimeSeriesMetadata::ChunkType; + if (parser.EatEqICase("UNCOMPRESSED")) { + create_option_.chunk_type = ChunkType::UNCOMPRESSED; + } else if (parser.EatEqICase("COMPRESSED")) { + create_option_.chunk_type = ChunkType::COMPRESSED; + } else { + return {Status::RedisParseErr, errBadEncoding}; + } + return Status::OK(); + } + + Status handleDuplicatePolicy(TSOptionsParser &parser) { + if (parser.EatEqICase("BLOCK")) { + create_option_.duplicate_policy = DuplicatePolicy::BLOCK; + } else if (parser.EatEqICase("FIRST")) { + create_option_.duplicate_policy = DuplicatePolicy::FIRST; + } else if (parser.EatEqICase("LAST")) { + create_option_.duplicate_policy = DuplicatePolicy::LAST; + } else if (parser.EatEqICase("MAX")) { + create_option_.duplicate_policy = DuplicatePolicy::MAX; + } else if (parser.EatEqICase("MIN")) { + create_option_.duplicate_policy = DuplicatePolicy::MIN; + } else if (parser.EatEqICase("SUM")) { + create_option_.duplicate_policy = DuplicatePolicy::SUM; + } else { + return {Status::RedisParseErr, errDuplicatePolicy}; + } + return Status::OK(); + } + + Status handleLabels(TSOptionsParser &parser) { + while (parser.Good()) { + auto parse_key = parser.TakeStr(); + auto parse_value = parser.TakeStr(); + if (!parse_key.IsOK() || !parse_value.IsOK()) { + break; + } + create_option_.labels.push_back({parse_key.GetValue(), parse_value.GetValue()}); + } + return Status::OK(); + } +}; + +class CommandTSCreate : public CommandTSCreateBase { + public: + CommandTSCreate() : CommandTSCreateBase(2, 0) {} + Status Parse(const std::vector &args) override { + if (args.size() < 2) { + return {Status::RedisParseErr, errWrongNumOfArguments}; + } + return CommandTSCreateBase::Parse(args); + } + Status Execute(engine::Context &ctx, Server *srv, Connection *conn, std::string *output) override { + auto timeseries_db = TimeSeries(srv->storage, conn->GetNamespace()); + auto s = timeseries_db.Create(ctx, args_[1], getCreateOption()); + if (!s.ok() && s.IsInvalidArgument()) return {Status::RedisExecErr, errKeyAlreadyExists}; + if (!s.ok()) return {Status::RedisExecErr, s.ToString()}; + *output = redis::RESP_OK; + return Status::OK(); + } +}; + +class CommandTSAdd : public CommandTSCreateBase { + public: + CommandTSAdd() : CommandTSCreateBase(4, 0) { + registerHandler("ON_DUPLICATE", [this](TSOptionsParser &parser) { return handleOnDuplicatePolicy(parser); }); + } + Status Parse(const std::vector &args) override { + if (args.size() < 4) { + return {Status::RedisParseErr, errWrongNumOfArguments}; + } + user_key_ = args[1]; + CommandParser parser(args, 2); + auto ts_parse = parser.TakeInt(); + if (!ts_parse.IsOK()) { + return {Status::RedisParseErr, errInvalidTimestamp}; + } + auto value_parse = parser.TakeFloat(); + if (!value_parse.IsOK()) { + return {Status::RedisParseErr, errInvalidValue}; + } + ts_ = ts_parse.GetValue(); + value_ = value_parse.GetValue(); + return CommandTSCreateBase::Parse(args); + } + Status Execute(engine::Context &ctx, Server *srv, Connection *conn, std::string *output) override { + auto timeseries_db = TimeSeries(srv->storage, conn->GetNamespace()); + const auto &option = getCreateOption(); + + TSChunk::AddResultWithTS res; + auto s = timeseries_db.Add(ctx, user_key_, {ts_, value_}, option, &res, + is_on_dup_policy_set_ ? &on_dup_policy_ : nullptr); + if (!s.ok()) return {Status::RedisExecErr, s.ToString()}; + + *output += FormatAddResultAsRedisReply(res); + + return Status::OK(); + } + + private: + DuplicatePolicy on_dup_policy_ = DuplicatePolicy::BLOCK; + bool is_on_dup_policy_set_ = false; + std::string user_key_; + uint64_t ts_ = 0; + double value_ = 0; + + Status handleOnDuplicatePolicy(TSOptionsParser &parser) { + if (parser.EatEqICase("BLOCK")) { + on_dup_policy_ = DuplicatePolicy::BLOCK; + } else if (parser.EatEqICase("FIRST")) { + on_dup_policy_ = DuplicatePolicy::FIRST; + } else if (parser.EatEqICase("LAST")) { + on_dup_policy_ = DuplicatePolicy::LAST; + } else if (parser.EatEqICase("MAX")) { + on_dup_policy_ = DuplicatePolicy::MAX; + } else if (parser.EatEqICase("MIN")) { + on_dup_policy_ = DuplicatePolicy::MIN; + } else if (parser.EatEqICase("SUM")) { + on_dup_policy_ = DuplicatePolicy::SUM; + } else { + return {Status::RedisParseErr, errDuplicatePolicy}; + } + is_on_dup_policy_set_ = true; + return Status::OK(); + } +}; + +class CommandTSMAdd : public Commander { + public: + Status Parse(const std::vector &args) override { + if (args.size() < 4 || (args.size() - 1) % 3 != 0) { + return {Status::RedisParseErr, errWrongNumOfArguments}; + } + CommandParser parser(args, 1); + for (size_t i = 1; i < args.size(); i += 3) { + const auto &user_key = args[i]; + parser.Skip(1); + auto ts_parse = parser.TakeInt(); + if (!ts_parse.IsOK()) { + return {Status::RedisParseErr, errInvalidTimestamp}; + } + auto value_parse = parser.TakeFloat(); + if (!value_parse.IsOK()) { + return Status{Status::RedisParseErr, errInvalidValue}; + } + userkey_samples_map_[user_key].push_back({ts_parse.GetValue(), value_parse.GetValue()}); + userkey_indexes_map_[user_key].push_back(i / 3); + samples_count_ += 1; + } + return Commander::Parse(args); + } + Status Execute(engine::Context &ctx, Server *srv, Connection *conn, std::string *output) override { + auto timeseries_db = TimeSeries(srv->storage, conn->GetNamespace()); + + auto replies = std::vector(samples_count_); + for (auto &[user_key, samples] : userkey_samples_map_) { + std::vector res; + auto count = samples.size(); + auto s = timeseries_db.MAdd(ctx, user_key, std::move(samples), &res); + std::string err_reply; + if (!s.ok()) { + err_reply = s.IsNotFound() ? redis::Error({Status::NotOK, errTSKeyNotFound}) + : redis::Error({Status::NotOK, s.ToString()}); + } + for (size_t i = 0; i < count; i++) { + size_t idx = userkey_indexes_map_[user_key][i]; + replies[idx] = s.ok() ? FormatAddResultAsRedisReply(res[i]) : err_reply; + } + } + *output = redis::MultiLen(samples_count_); + for (auto &reply : replies) { + if (reply.empty()) continue; + *output += reply; + } + return Status::OK(); + } + + private: + std::string user_key_; + size_t samples_count_ = 0; + std::unordered_map> userkey_samples_map_; + std::unordered_map> userkey_indexes_map_; +}; + +REDIS_REGISTER_COMMANDS(Timeseries, MakeCmdAttr("ts.create", -2, "write", 1, 1, 1), + MakeCmdAttr("ts.add", -4, "write", 1, 1, 1), + MakeCmdAttr("ts.madd", -4, "write", 1, -3, 1), ); + +} // namespace redis diff --git a/src/commands/commander.h b/src/commands/commander.h index 8b012ff27f1..a44f68505ee 100644 --- a/src/commands/commander.h +++ b/src/commands/commander.h @@ -111,6 +111,7 @@ enum class CommandCategory : uint8_t { TDigest, Txn, ZSet, + Timeseries, }; class Commander { diff --git a/src/types/redis_timeseries.cc b/src/types/redis_timeseries.cc index 73023c35c87..55c299a78f6 100644 --- a/src/types/redis_timeseries.cc +++ b/src/types/redis_timeseries.cc @@ -20,8 +20,17 @@ #include "redis_timeseries.h" +#include "commands/error_constants.h" +#include "db_util.h" + namespace redis { +// TODO: make it configurable +constexpr uint64_t kDefaultRetentionTime = 0; +constexpr uint64_t kDefaultChunkSize = 1024; +constexpr auto kDefaultChunkType = TimeSeriesMetadata::ChunkType::UNCOMPRESSED; +constexpr auto kDefaultDuplicatePolicy = TimeSeriesMetadata::DuplicatePolicy::BLOCK; + void TSDownStreamMeta::Encode(std::string *dst) const { PutFixed8(dst, static_cast(aggregator)); PutFixed64(dst, bucket_duration); @@ -88,7 +97,185 @@ std::string TSRevLabelKey::Encode() const { return encoded; } -std::string TimeSeries::internalKeyFromChunkID(const std::string &ns_key, const TimeSeriesMetadata &metadata, +TSCreateOption::TSCreateOption() + : retention_time(kDefaultRetentionTime), + chunk_size(kDefaultChunkSize), + chunk_type(kDefaultChunkType), + duplicate_policy(kDefaultDuplicatePolicy) {} + +TimeSeriesMetadata CreateMetadataFromOption(const TSCreateOption &option) { + TimeSeriesMetadata metadata; + metadata.retention_time = option.retention_time; + metadata.chunk_size = option.chunk_size; + metadata.chunk_type = option.chunk_type; + metadata.duplicate_policy = option.duplicate_policy; + metadata.SetSourceKey(option.source_key); + + return metadata; +} + +rocksdb::Status TimeSeries::getTimeSeriesMetadata(engine::Context &ctx, const Slice &ns_key, + TimeSeriesMetadata *metadata) { + return Database::GetMetadata(ctx, {kRedisTimeSeries}, ns_key, metadata); +} + +rocksdb::Status TimeSeries::createTimeSeries(engine::Context &ctx, const Slice &ns_key, + TimeSeriesMetadata *metadata_out, const TSCreateOption *option) { + auto batch = storage_->GetWriteBatchBase(); + WriteBatchLogData log_data(kRedisTimeSeries, {"createTimeSeries"}); + auto s = batch->PutLogData(log_data.Encode()); + if (!s.ok()) return s; + + *metadata_out = CreateMetadataFromOption(option ? *option : TSCreateOption{}); + std::string bytes; + metadata_out->Encode(&bytes); + s = batch->Put(metadata_cf_handle_, ns_key, bytes); + if (!s.ok()) return s; + + if (!option && !option->labels.empty()) { + createLabelIndexInBatch(ns_key, *metadata_out, batch, option->labels); + } + + return storage_->Write(ctx, storage_->DefaultWriteOptions(), batch->GetWriteBatch()); +} + +rocksdb::Status TimeSeries::getOrCreateTimeSeries(engine::Context &ctx, const Slice &ns_key, + TimeSeriesMetadata *metadata_out, const TSCreateOption *option) { + auto s = getTimeSeriesMetadata(ctx, ns_key, metadata_out); + if (s.ok()) { + return s; + } + return createTimeSeries(ctx, ns_key, metadata_out, option); +} + +rocksdb::Status TimeSeries::upsertCommon(engine::Context &ctx, const Slice &ns_key, TimeSeriesMetadata &metadata, + SampleBatch &sample_batch) { + auto all_batch_slice = sample_batch.AsSlice(); + + // In the emun `TSSubkeyType`, `LABEL` is the next of `CHUNK` + std::string chunk_upper_bound = internalKeyFromLabelKey(ns_key, metadata, ""); + std::string end_key = internalKeyFromChunkID(ns_key, metadata, TSSample::MAX_TIMESTAMP); + std::string prefix = end_key.substr(0, end_key.size() - sizeof(uint64_t)); + + rocksdb::ReadOptions read_options = ctx.DefaultScanOptions(); + rocksdb::Slice upper_bound(chunk_upper_bound); + read_options.iterate_upper_bound = &upper_bound; + rocksdb::Slice lower_bound(prefix); + read_options.iterate_lower_bound = &lower_bound; + + uint64_t chunk_count = metadata.size; + + // Get the latest chunk + auto iter = util::UniqueIterator(ctx, read_options); + iter->SeekForPrev(end_key); + TSChunkPtr latest_chunk; + std::string latest_chunk_key, latest_chunk_value; + if (!iter->Valid() || !iter->key().starts_with(prefix)) { + // Create a new empty chunk if there is no chunk + auto [chunk_ptr_, data_] = CreateEmptyOwnedTSChunk(); + latest_chunk_value = std::move(data_); + latest_chunk = std::move(chunk_ptr_); + } else { + latest_chunk_key = iter->key().ToString(); + latest_chunk_value = iter->value().ToString(); + latest_chunk = CreateTSChunkFromData(latest_chunk_value); + } + + // Filter out samples older than retention time + sample_batch.Expire(latest_chunk->GetLastTimestamp(), metadata.retention_time); + if (all_batch_slice.GetValidCount() == 0) { + return rocksdb::Status::OK(); + } + + auto batch = storage_->GetWriteBatchBase(); + WriteBatchLogData log_data(kRedisTimeSeries); + auto s = batch->PutLogData(log_data.Encode()); + if (!s.ok()) return s; + + // Get the first chunk + auto start_key = internalKeyFromChunkID(ns_key, metadata, all_batch_slice.GetFirstTimestamp()); + iter->SeekForPrev(start_key); + if (!iter->Valid()) { + iter->Seek(start_key); + } else if (!iter->key().starts_with(prefix)) { + iter->Next(); + } + + // Process samples added to sealed chunks + uint64_t start_ts = 0; + uint64_t end_ts = TSSample::MAX_TIMESTAMP; + bool is_chunk = (iter->Valid() && iter->key().starts_with(prefix)); + while (is_chunk) { + auto cur_chunk_data = iter->value().ToString(); + auto cur_chunk_key = iter->key().ToString(); + iter->Next(); + is_chunk = (iter->Valid() && iter->key().starts_with(prefix)); + if (!is_chunk) { + // Process last chunk + break; + } + end_ts = chunkIDFromInternalKey(iter->key()); + + auto chunk = CreateTSChunkFromData(cur_chunk_data); + auto sample_slice = all_batch_slice.SliceByTimestamps(start_ts, end_ts); + if (sample_slice.GetValidCount() == 0) { + continue; + } + auto new_data_list = chunk->UpsertSampleAndSplit(sample_slice, metadata.chunk_size, false); + for (size_t i = 0; i < new_data_list.size(); i++) { + auto &new_data = new_data_list[i]; + auto new_chunk = CreateTSChunkFromData(new_data); + auto new_key = internalKeyFromChunkID(ns_key, metadata, new_chunk->GetFirstTimestamp()); + // Process samples older than the first chunk, should update the key + if (i == 0 && new_key != cur_chunk_key) { + s = batch->Delete(cur_chunk_key); + if (!s.ok()) return s; + } + s = batch->Put(new_key, new_data); + if (!s.ok()) return s; + } + chunk_count += new_data_list.size() - 1; + } + + // Process samples added to latest chunk(unseal) + auto remained_samples = all_batch_slice.SliceByTimestamps(start_ts, TSSample::MAX_TIMESTAMP, true); + + auto new_data_list = latest_chunk->UpsertSampleAndSplit(remained_samples, metadata.chunk_size, true); + for (size_t i = 0; i < new_data_list.size(); i++) { + auto &new_data = new_data_list[i]; + auto new_chunk = CreateTSChunkFromData(new_data); + auto new_key = internalKeyFromChunkID(ns_key, metadata, new_chunk->GetFirstTimestamp()); + if (i == 0 && new_key != latest_chunk_key) { + s = batch->Delete(latest_chunk_key); + if (!s.ok()) return s; + } + s = batch->Put(new_key, new_data); + if (!s.ok()) return s; + } + chunk_count += new_data_list.size() - (metadata.size == 0 ? 0 : 1); + if (chunk_count != metadata.size) { + metadata.size = chunk_count; + std::string bytes; + metadata.Encode(&bytes); + s = batch->Put(metadata_cf_handle_, ns_key, bytes); + if (!s.ok()) return s; + } + + return storage_->Write(ctx, storage_->DefaultWriteOptions(), batch->GetWriteBatch()); +} + +rocksdb::Status TimeSeries::createLabelIndexInBatch(const Slice &ns_key, const TimeSeriesMetadata &metadata, + ObserverOrUniquePtr &batch, + const LabelKVList &labels) { + for (auto &label : labels) { + auto internal_key = internalKeyFromLabelKey(ns_key, metadata, label.k); + auto s = batch->Put(internal_key, label.v); + if (!s.ok()) return s; + } + return rocksdb::Status::OK(); +} + +std::string TimeSeries::internalKeyFromChunkID(const Slice &ns_key, const TimeSeriesMetadata &metadata, uint64_t id) const { std::string sub_key; PutFixed8(&sub_key, static_cast(TSSubkeyType::CHUNK)); @@ -97,7 +284,7 @@ std::string TimeSeries::internalKeyFromChunkID(const std::string &ns_key, const return InternalKey(ns_key, sub_key, metadata.version, storage_->IsSlotIdEncoded()).Encode(); } -std::string TimeSeries::internalKeyFromLabelKey(const std::string &ns_key, const TimeSeriesMetadata &metadata, +std::string TimeSeries::internalKeyFromLabelKey(const Slice &ns_key, const TimeSeriesMetadata &metadata, Slice label_key) const { std::string sub_key; sub_key.resize(1 + label_key.size()); @@ -108,7 +295,7 @@ std::string TimeSeries::internalKeyFromLabelKey(const std::string &ns_key, const return InternalKey(ns_key, sub_key, metadata.version, storage_->IsSlotIdEncoded()).Encode(); } -std::string TimeSeries::internalKeyFromDownstreamKey(const std::string &ns_key, const TimeSeriesMetadata &metadata, +std::string TimeSeries::internalKeyFromDownstreamKey(const Slice &ns_key, const TimeSeriesMetadata &metadata, Slice downstream_key) const { std::string sub_key; sub_key.resize(1 + downstream_key.size()); @@ -119,4 +306,57 @@ std::string TimeSeries::internalKeyFromDownstreamKey(const std::string &ns_key, return InternalKey(ns_key, sub_key, metadata.version, storage_->IsSlotIdEncoded()).Encode(); } +uint64_t TimeSeries::chunkIDFromInternalKey(Slice internal_key) { + auto size = internal_key.size(); + internal_key.remove_prefix(size - sizeof(uint64_t)); + return DecodeFixed64(internal_key.data()); +} + +rocksdb::Status TimeSeries::Create(engine::Context &ctx, const Slice &user_key, const TSCreateOption &option) { + std::string ns_key = AppendNamespacePrefix(user_key); + + TimeSeriesMetadata metadata; + auto s = getTimeSeriesMetadata(ctx, ns_key, &metadata); + if (s.ok()) { + return rocksdb::Status::InvalidArgument("key already exists"); + } + return createTimeSeries(ctx, ns_key, &metadata, &option); +} + +rocksdb::Status TimeSeries::Add(engine::Context &ctx, const Slice &user_key, TSSample sample, + const TSCreateOption &option, AddResultWithTS *res, + const DuplicatePolicy *on_dup_policy) { + std::string ns_key = AppendNamespacePrefix(user_key); + + TimeSeriesMetadata metadata(false); + rocksdb::Status s = getOrCreateTimeSeries(ctx, ns_key, &metadata, &option); + if (!s.ok()) return s; + auto sample_batch = SampleBatch({sample}, on_dup_policy ? *on_dup_policy : metadata.duplicate_policy); + + s = upsertCommon(ctx, ns_key, metadata, sample_batch); + if (!s.ok()) { + return s; + } + *res = sample_batch.GetFinalResults()[0]; + return rocksdb::Status::OK(); +} + +rocksdb::Status TimeSeries::MAdd(engine::Context &ctx, const Slice &user_key, std::vector samples, + std::vector *res) { + std::string ns_key = AppendNamespacePrefix(user_key); + + TimeSeriesMetadata metadata(false); + rocksdb::Status s = getTimeSeriesMetadata(ctx, ns_key, &metadata); + if (!s.ok()) { + return s; + } + auto sample_batch = SampleBatch(std::move(samples), metadata.duplicate_policy); + s = upsertCommon(ctx, ns_key, metadata, sample_batch); + if (!s.ok()) { + return s; + } + *res = sample_batch.GetFinalResults(); + return rocksdb::Status::OK(); +} + } // namespace redis diff --git a/src/types/redis_timeseries.h b/src/types/redis_timeseries.h index e28b1120024..c853ec83363 100644 --- a/src/types/redis_timeseries.h +++ b/src/types/redis_timeseries.h @@ -24,6 +24,7 @@ #include "storage/redis_db.h" #include "storage/redis_metadata.h" +#include "types/timeseries.h" namespace redis { @@ -88,16 +89,54 @@ struct TSRevLabelKey { [[nodiscard]] std::string Encode() const; }; +struct LabelKVPair { + std::string k; + std::string v; +}; +using LabelKVList = std::vector; + +struct TSCreateOption { + uint64_t retention_time; + uint64_t chunk_size; + TimeSeriesMetadata::ChunkType chunk_type; + TimeSeriesMetadata::DuplicatePolicy duplicate_policy; + std::string source_key; + LabelKVList labels; + + TSCreateOption(); +}; + +TimeSeriesMetadata CreateMetadataFromOption(const TSCreateOption &option); + class TimeSeries : public SubKeyScanner { public: + using SampleBatch = TSChunk::SampleBatch; + using AddResultWithTS = TSChunk::AddResultWithTS; + using DuplicatePolicy = TimeSeriesMetadata::DuplicatePolicy; + TimeSeries(engine::Storage *storage, const std::string &ns) : SubKeyScanner(storage, ns) {} + rocksdb::Status Create(engine::Context &ctx, const Slice &user_key, const TSCreateOption &option); + rocksdb::Status Add(engine::Context &ctx, const Slice &user_key, TSSample sample, const TSCreateOption &option, + AddResultWithTS *res, const DuplicatePolicy *on_dup_policy = nullptr); + rocksdb::Status MAdd(engine::Context &ctx, const Slice &user_key, std::vector samples, + std::vector *res); private: - std::string internalKeyFromChunkID(const std::string &ns_key, const TimeSeriesMetadata &metadata, uint64_t id) const; - std::string internalKeyFromLabelKey(const std::string &ns_key, const TimeSeriesMetadata &metadata, - Slice label_key) const; - std::string internalKeyFromDownstreamKey(const std::string &ns_key, const TimeSeriesMetadata &metadata, + rocksdb::Status getTimeSeriesMetadata(engine::Context &ctx, const Slice &ns_key, TimeSeriesMetadata *metadata); + rocksdb::Status createTimeSeries(engine::Context &ctx, const Slice &ns_key, TimeSeriesMetadata *metadata_out, + const TSCreateOption *options); + rocksdb::Status getOrCreateTimeSeries(engine::Context &ctx, const Slice &ns_key, TimeSeriesMetadata *metadata_out, + const TSCreateOption *option = nullptr); + rocksdb::Status upsertCommon(engine::Context &ctx, const Slice &ns_key, TimeSeriesMetadata &metadata, + SampleBatch &sample_batch); + rocksdb::Status createLabelIndexInBatch(const Slice &ns_key, const TimeSeriesMetadata &metadata, + ObserverOrUniquePtr &batch, + const LabelKVList &labels); + std::string internalKeyFromChunkID(const Slice &ns_key, const TimeSeriesMetadata &metadata, uint64_t id) const; + std::string internalKeyFromLabelKey(const Slice &ns_key, const TimeSeriesMetadata &metadata, Slice label_key) const; + std::string internalKeyFromDownstreamKey(const Slice &ns_key, const TimeSeriesMetadata &metadata, Slice downstream_key) const; + static uint64_t chunkIDFromInternalKey(Slice internal_key); }; } // namespace redis diff --git a/src/types/timeseries.cc b/src/types/timeseries.cc index 89f0e1a8c52..5fb505dcc73 100644 --- a/src/types/timeseries.cc +++ b/src/types/timeseries.cc @@ -28,12 +28,12 @@ using AddResult = TSChunk::AddResult; using SampleBatch = TSChunk::SampleBatch; using SampleBatchSlice = TSChunk::SampleBatchSlice; -TSChunkPtr CreateTSChunkFromData(nonstd::span data) { +TSChunkPtr CreateTSChunkFromData(nonstd::span data) { auto chunk_meta = TSChunk::MetaData(); Slice input(data.data(), TSChunk::MetaData::kEncodedSize); chunk_meta.Decode(&input); if (!chunk_meta.is_compressed) { - return std::make_unique(std::move(data)); + return std::make_unique(data); } else { // TODO: compressed chunk unreachable(); @@ -59,7 +59,11 @@ TSChunk::SampleBatch::SampleBatch(std::vector samples, DuplicatePolicy void TSChunk::SampleBatch::Expire(uint64_t last_ts, uint64_t retention) { if (retention == 0) return; - for (auto idx : indexes_) { + std::vector inverse(indexes_.size()); + for (size_t i = 0; i < indexes_.size(); ++i) { + inverse[indexes_[i]] = i; + } + for (auto idx : inverse) { if (samples_[idx].ts + retention < last_ts) { add_results_[idx] = AddResult::kOld; } else if (samples_[idx].ts > last_ts) { @@ -144,7 +148,7 @@ SampleBatchSlice TSChunk::SampleBatchSlice::SliceByTimestamps(uint64_t first, ui return {}; } -SampleBatchSlice TSChunk::SampleBatchSlice::createSampleSlice(size_t start_idx, size_t end_idx) { +SampleBatchSlice TSChunk::SampleBatchSlice::createSampleSlice(size_t start_idx, size_t end_idx) const { if (end_idx > sample_span_.size()) { end_idx = sample_span_.size(); } @@ -157,11 +161,12 @@ SampleBatchSlice TSChunk::SampleBatchSlice::createSampleSlice(size_t start_idx, SampleBatchSlice TSChunk::SampleBatch::AsSlice() { return {samples_, add_results_, policy_}; } -std::vector TSChunk::SampleBatch::GetFinalResults() const { - std::vector res; +std::vector TSChunk::SampleBatch::GetFinalResults() const { + std::vector res; res.resize(add_results_.size()); for (size_t idx = 0; idx < add_results_.size(); idx++) { - res[indexes_[idx]] = add_results_[idx]; + res[indexes_[idx]].first = add_results_[idx]; + res[indexes_[idx]].second = samples_[idx].ts; } return res; } @@ -195,7 +200,7 @@ AddResult TSChunk::MergeSamplesValue(TSSample& to, const TSSample& from, Duplica uint32_t TSChunk::GetCount() const { return metadata_.count; } -uint64_t TSChunk::SampleBatchSlice::GetFirstTimestamp() { +uint64_t TSChunk::SampleBatchSlice::GetFirstTimestamp() const { if (sample_span_.size() == 0) return 0; for (size_t i = 0; i < sample_span_.size(); i++) { if (add_result_span_[i] == AddResult::kNone) { @@ -205,11 +210,12 @@ uint64_t TSChunk::SampleBatchSlice::GetFirstTimestamp() { return 0; } -uint64_t TSChunk::SampleBatchSlice::GetLastTimestamp() { +uint64_t TSChunk::SampleBatchSlice::GetLastTimestamp() const { if (sample_span_.size() == 0) return 0; - for (size_t i = sample_span_.size() - 1; i >= 0; i--) { - if (add_result_span_[i] == AddResult::kNone) { - return sample_span_[i].ts; + for (size_t i = 0; i < sample_span_.size(); i++) { + auto index = sample_span_.size() - i - 1; + if (add_result_span_[index] == AddResult::kNone) { + return sample_span_[index].ts; } } return 0; @@ -242,26 +248,27 @@ void TSChunk::MetaData::Decode(Slice* input) { GetFixed32(input, &count); } -TSChunk::TSChunk(nonstd::span data) : data_(data) { +TSChunk::TSChunk(nonstd::span data) : data_(data) { Slice input(data_.data(), data_.size()); metadata_.Decode(&input); } class UncompTSChunkIterator : public TSChunkIterator { public: - explicit UncompTSChunkIterator(nonstd::span data, uint64_t count) : TSChunkIterator(count), data_(data) {} - std::optional Next() override { + explicit UncompTSChunkIterator(nonstd::span data, uint64_t count) + : TSChunkIterator(count), data_(data) {} + std::optional Next() override { if (idx_ >= count_) return std::nullopt; return &data_[idx_++]; } private: - nonstd::span data_; + nonstd::span data_; }; -UncompTSChunk::UncompTSChunk(nonstd::span data) : TSChunk(data) { - auto data_ptr = reinterpret_cast(data.data()) + TSChunk::MetaData::kEncodedSize; - samples_ = nonstd::span(reinterpret_cast(data_ptr), metadata_.count); +UncompTSChunk::UncompTSChunk(nonstd::span data) : TSChunk(data) { + auto data_ptr = reinterpret_cast(data.data()) + TSChunk::MetaData::kEncodedSize; + samples_ = nonstd::span(reinterpret_cast(data_ptr), metadata_.count); } std::unique_ptr UncompTSChunk::CreateIterator() const { @@ -323,7 +330,7 @@ std::string UncompTSChunk::UpsertSamples(SampleBatchSlice batch) const { // Select next sample by earliest timestamp if (existing_sample_iter->ts <= new_samples[new_sample_idx].ts) { - candidate = &(*existing_sample_iter); + candidate = existing_sample_iter; } else { candidate = &new_samples[new_sample_idx]; from_new_batch = true; @@ -337,17 +344,21 @@ std::string UncompTSChunk::UpsertSamples(SampleBatchSlice batch) const { if (current_index == static_cast(-1)) { merged_data[0] = *candidate; current_index = 0; + if (from_new_batch) { + add_results[new_sample_idx] = AddResult::kOk; + } continue; } // Append or merge based on timestamp + bool is_append = false; if (candidate->ts > merged_data[current_index].ts) { merged_data[++current_index] = *candidate; - } else { - if (from_new_batch) { - auto add_res = MergeSamplesValue(merged_data[current_index], *candidate, policy); - add_results[new_sample_idx] = add_res; - } + is_append = true; + } + if (from_new_batch) { + add_results[new_sample_idx] = + is_append ? AddResult::kOk : MergeSamplesValue(merged_data[current_index], *candidate, policy); } // Update the index @@ -361,7 +372,7 @@ std::string UncompTSChunk::UpsertSamples(SampleBatchSlice batch) const { // Copy remaining existing samples if (existing_sample_iter != samples_.end()) { const size_t remaining_count = std::distance(existing_sample_iter, samples_.end()); - std::memcpy(&merged_data[current_index + 1], &(*existing_sample_iter), remaining_count * sizeof(TSSample)); + std::memcpy(&merged_data[current_index + 1], existing_sample_iter, remaining_count * sizeof(TSSample)); current_index += remaining_count; } @@ -374,8 +385,10 @@ std::string UncompTSChunk::UpsertSamples(SampleBatchSlice batch) const { if (current_index == static_cast(-1)) { current_index = 0; merged_data[current_index] = new_samples[new_sample_idx]; + add_results[new_sample_idx] = AddResult::kOk; } else if (new_samples[new_sample_idx].ts > merged_data[current_index].ts) { merged_data[++current_index] = new_samples[new_sample_idx]; + add_results[new_sample_idx] = AddResult::kOk; } else { auto add_res = MergeSamplesValue(merged_data[current_index], new_samples[new_sample_idx], policy); add_results[new_sample_idx] = add_res; @@ -393,6 +406,63 @@ std::string UncompTSChunk::UpsertSamples(SampleBatchSlice batch) const { return new_buffer; } +std::vector UncompTSChunk::UpsertSampleAndSplit(SampleBatchSlice batch, uint64_t preferred_chunk_size, + bool is_fix_split_mode) const { + auto whole_chunk_data = UpsertSamples(batch); + // Return empty if no changes + if (whole_chunk_data.empty()) { + return {}; + } + auto whole_chunk = CreateTSChunkFromData(whole_chunk_data); + + // Split + std::vector split_size; + auto total_count = whole_chunk->GetCount(); + if (is_fix_split_mode) { + // Fixed split + size_t remaining = total_count; + while (remaining > 0) { + auto size = std::min(remaining, preferred_chunk_size); + split_size.push_back(size); + remaining -= size; + } + } else if (total_count > 2 * preferred_chunk_size) { + // Equal split + auto split_count = total_count / preferred_chunk_size; + auto chunk_size = total_count / split_count; + auto remainder = total_count % split_count; + split_size.resize(split_count); + std::fill(split_size.begin(), split_size.end(), chunk_size); + for (uint32_t i = 0; i < remainder; ++i) { + split_size[i] += 1; + } + } + if (split_size.empty()) { + split_size.push_back(total_count); + } + // Return if only one chunk + if (split_size.size() == 1) { + return {std::move(whole_chunk_data)}; + } + + constexpr size_t header_size = TSChunk::MetaData::kEncodedSize; + const char* data_ptr = whole_chunk_data.data() + header_size; + std::vector res; + for (auto size : split_size) { + auto sample_bytes = size * sizeof(TSSample); + const size_t required_size = header_size + sample_bytes; + std::string buffer; + buffer.resize(required_size); + auto metadata = TSChunk::MetaData(false, size); + auto str = metadata.Encode(); + EncodeBuffer(buffer.data(), str); + std::memcpy(buffer.data() + header_size, data_ptr, sample_bytes); + data_ptr += sample_bytes; + res.push_back(std::move(buffer)); + } + return res; +} + std::string UncompTSChunk::RemoveSamplesBetween(uint64_t from, uint64_t to) const { if (from > to) { return ""; diff --git a/src/types/timeseries.h b/src/types/timeseries.h index e5f82ebec94..f0e883ad28c 100644 --- a/src/types/timeseries.h +++ b/src/types/timeseries.h @@ -32,7 +32,7 @@ using TSChunkPtr = std::unique_ptr; using OwnedTSChunk = std::tuple; // Creates a TSChunk from the provided raw data buffer. -TSChunkPtr CreateTSChunkFromData(nonstd::span data); +TSChunkPtr CreateTSChunkFromData(nonstd::span data); // Creates an empty owned time series chunk with specified compression option. OwnedTSChunk CreateEmptyOwnedTSChunk(bool is_compressed = false); @@ -54,7 +54,7 @@ class TSChunkIterator { explicit TSChunkIterator(uint64_t count) : count_(count), idx_(0) {} virtual ~TSChunkIterator() = default; - virtual std::optional Next() = 0; + virtual std::optional Next() = 0; virtual bool HasNext() const { return idx_ < count_; } protected: @@ -72,6 +72,7 @@ class TSChunk { kBlock, kOld, }; + using AddResultWithTS = std::pair; class SampleBatch; class SampleBatchSlice { @@ -88,8 +89,8 @@ class TSChunk { // e.g., samples: {10,20,30,40}, first=20, last=40 -> {20,30} SampleBatchSlice SliceByTimestamps(uint64_t first, uint64_t last, bool contain_last = false); - uint64_t GetFirstTimestamp(); - uint64_t GetLastTimestamp(); + uint64_t GetFirstTimestamp() const; + uint64_t GetLastTimestamp() const; // Get number of valid samples (excluding duplicates and expired entries) size_t GetValidCount() const; @@ -109,7 +110,7 @@ class TSChunk { SampleBatchSlice(nonstd::span samples, nonstd::span results, DuplicatePolicy policy) : sample_span_(samples), add_result_span_(results), policy_(policy) {} - SampleBatchSlice createSampleSlice(size_t start_idx, size_t end_idx); + SampleBatchSlice createSampleSlice(size_t start_idx, size_t end_idx) const; }; class SampleBatch { @@ -125,7 +126,7 @@ class TSChunk { SampleBatchSlice AsSlice(); // Return add results by samples' order - std::vector GetFinalResults() const; + std::vector GetFinalResults() const; private: std::vector samples_; @@ -148,7 +149,7 @@ class TSChunk { void Decode(Slice* input); }; - explicit TSChunk(nonstd::span data); + explicit TSChunk(nonstd::span data); virtual ~TSChunk() = default; @@ -166,6 +167,13 @@ class TSChunk { // Returns new chunk data with merged samples. Returns empty string if no changes virtual std::string UpsertSamples(SampleBatchSlice samples) const = 0; + // Add new samples to the chunk according to duplicate policy + // Split chunk and return new chunk. There two split modes: + // 1. Fix split mode: used for unsealed chunk. 2. Equal split mode: used for sealed chunk. + // Returns empty if no changes + virtual std::vector UpsertSampleAndSplit(SampleBatchSlice batch, uint64_t preferred_chunk_size, + bool is_fix_split_mode) const = 0; + // Delete samples in [from, to] timestamp range // Returns new chunk data without deleted samples. Returns empty string if no changes virtual std::string RemoveSamplesBetween(uint64_t from, uint64_t to) const = 0; @@ -176,22 +184,24 @@ class TSChunk { virtual std::string UpdateSampleValue(uint64_t ts, double value, bool is_add_on) const = 0; protected: - nonstd::span data_; + nonstd::span data_; MetaData metadata_; }; class UncompTSChunk : public TSChunk { public: - explicit UncompTSChunk(nonstd::span data); + explicit UncompTSChunk(nonstd::span data); std::unique_ptr CreateIterator() const override; uint64_t GetFirstTimestamp() const override; uint64_t GetLastTimestamp() const override; std::string UpsertSamples(SampleBatchSlice samples) const override; + std::vector UpsertSampleAndSplit(SampleBatchSlice batch, uint64_t preferred_chunk_size, + bool is_fix_split_mode) const override; std::string RemoveSamplesBetween(uint64_t from, uint64_t to) const override; std::string UpdateSampleValue(uint64_t ts, double value, bool is_add_on) const override; private: - nonstd::span samples_; + nonstd::span samples_; }; diff --git a/tests/cppunit/types/timeseries_chunk_test.cc b/tests/cppunit/types/timeseries_chunk_test.cc index dbbf0912e89..c8d89c13e6d 100644 --- a/tests/cppunit/types/timeseries_chunk_test.cc +++ b/tests/cppunit/types/timeseries_chunk_test.cc @@ -108,11 +108,12 @@ TEST(RedisTimeSeriesChunkTest, ExpirationLogic) { batch.Expire(300, 150); auto results = batch.GetFinalResults(); - // Only samples with ts >= 150 should be kept - EXPECT_EQ(results[0], AddResult::kNone); - EXPECT_EQ(results[1], AddResult::kNone); - EXPECT_EQ(results[2], AddResult::kOld); - EXPECT_EQ(results[3], AddResult::kOld); + EXPECT_EQ(results[0].first, AddResult::kNone); + EXPECT_EQ(results[0].second, 200); + EXPECT_EQ(results[1].first, AddResult::kNone); + EXPECT_EQ(results[1].second, 400); + EXPECT_EQ(results[2].first, AddResult::kOld); + EXPECT_EQ(results[3].first, AddResult::kOld); } // Test SampleBatch construction and sorting @@ -133,11 +134,14 @@ TEST(RedisTimeSeriesChunkTest, BatchSortingAndDeduplication) { // Verify deduplication EXPECT_EQ(slice.GetValidCount(), 3); auto results = batch.GetFinalResults(); - EXPECT_EQ(results[0], AddResult::kNone); - EXPECT_EQ(results[1], AddResult::kNone); - EXPECT_EQ(results[2], AddResult::kNone); - EXPECT_EQ(results[3], AddResult::kBlock); - EXPECT_EQ(results[4], AddResult::kBlock); + EXPECT_EQ(results[0].first, AddResult::kNone); + EXPECT_EQ(results[0].second, 300); + EXPECT_EQ(results[1].first, AddResult::kNone); + EXPECT_EQ(results[1].second, 100); + EXPECT_EQ(results[2].first, AddResult::kNone); + EXPECT_EQ(results[2].second, 200); + EXPECT_EQ(results[3].first, AddResult::kBlock); + EXPECT_EQ(results[4].first, AddResult::kBlock); } // Test MAddSample merging logic with additional samples and content validation @@ -163,6 +167,23 @@ TEST(RedisTimeSeriesChunkTest, UcompChunkMAddSampleLogic) { EXPECT_EQ(new_chunk->GetFirstTimestamp(), 100); EXPECT_EQ(new_chunk->GetLastTimestamp(), 400); + // Verify add result + auto results = batch.GetFinalResults(); + EXPECT_EQ(results[0].first, AddResult::kOk); + EXPECT_EQ(results[0].second, 300); + EXPECT_EQ(results[1].first, AddResult::kOk); + EXPECT_EQ(results[1].second, 100); + EXPECT_EQ(results[2].first, AddResult::kOk); + EXPECT_EQ(results[2].second, 200); + EXPECT_EQ(results[3].first, AddResult::kOk); + EXPECT_EQ(results[3].second, 100); + EXPECT_EQ(results[4].first, AddResult::kOk); + EXPECT_EQ(results[4].second, 200); + EXPECT_EQ(results[5].first, AddResult::kOk); + EXPECT_EQ(results[5].second, 400); + EXPECT_EQ(results[6].first, AddResult::kOk); + EXPECT_EQ(results[6].second, 100); + // Validate content of merged chunk auto iter = new_chunk->CreateIterator(); auto* sample = iter->Next().value(); @@ -209,6 +230,19 @@ TEST(RedisTimeSeriesChunkTest, UcompChunkMAddSampleWithExistingSamples) { EXPECT_EQ(final_chunk->GetFirstTimestamp(), 50); EXPECT_EQ(final_chunk->GetLastTimestamp(), 400); + // Verify add result + auto results = new_batch.GetFinalResults(); + EXPECT_EQ(results[0].first, AddResult::kOk); + EXPECT_EQ(results[0].second, 50); + EXPECT_EQ(results[1].first, AddResult::kOk); + EXPECT_EQ(results[1].second, 150); + EXPECT_EQ(results[2].first, AddResult::kOk); + EXPECT_EQ(results[2].second, 200); + EXPECT_EQ(results[3].first, AddResult::kOk); + EXPECT_EQ(results[3].second, 300); + EXPECT_EQ(results[4].first, AddResult::kOk); + EXPECT_EQ(results[4].second, 400); + // Verify content through iterator auto iter = final_chunk->CreateIterator(); auto* sample = iter->Next().value(); @@ -372,4 +406,83 @@ TEST(RedisTimeSeriesChunkTest, UpdateSampleBehavior) { EXPECT_TRUE(updated_data.empty()); } +// Test UpsertSampleAndSplit with different split modes and chunk size requirements +TEST(RedisTimeSeriesChunkTest, UcompChunkUpsertAndSplitBehaviors) { + // Base chunk with 3 samples + auto [chunk, data] = CreateEmptyOwnedTSChunk(false); + std::vector base_samples = {MakeSample(100, 1.0), MakeSample(200, 2.0), MakeSample(300, 3.0)}; + SampleBatch base_batch(base_samples, DuplicatePolicy::LAST); + std::string merged_data = chunk->UpsertSamples(base_batch.AsSlice()); + ASSERT_FALSE(merged_data.empty()); + + // Test case 1: No split needed (chunk size exactly matches preferred) + auto test_chunk = CreateTSChunkFromData(merged_data); + std::vector new_samples = {MakeSample(400, 4.0)}; + SampleBatch new_batch(new_samples, DuplicatePolicy::LAST); + auto result = test_chunk->UpsertSampleAndSplit(new_batch.AsSlice(), 4, false); + ASSERT_EQ(result.size(), 1); + auto result_chunk = CreateTSChunkFromData(result[0]); + EXPECT_EQ(result_chunk->GetCount(), 4); + EXPECT_EQ(result_chunk->GetFirstTimestamp(), 100); + EXPECT_EQ(result_chunk->GetLastTimestamp(), 400); + + // Test case 2: Fixed split mode (7 samples into 3 chunks of 3,3 and 1) + test_chunk = CreateTSChunkFromData(merged_data); + new_samples = {MakeSample(400, 4.0), MakeSample(500, 5.0), MakeSample(600, 6.0), MakeSample(700, 7.0)}; + new_batch = SampleBatch(new_samples, DuplicatePolicy::LAST); + result = test_chunk->UpsertSampleAndSplit(new_batch.AsSlice(), 3, true); + ASSERT_EQ(result.size(), 3); + EXPECT_EQ(result[0].size(), TSChunk::MetaData::kEncodedSize + 3 * sizeof(TSSample)); + EXPECT_EQ(result[1].size(), TSChunk::MetaData::kEncodedSize + 3 * sizeof(TSSample)); + EXPECT_EQ(result[2].size(), TSChunk::MetaData::kEncodedSize + 1 * sizeof(TSSample)); + + // Validate first chunk content + auto chunk1 = CreateTSChunkFromData(result[0]); + EXPECT_EQ(chunk1->GetCount(), 3); + auto iter = chunk1->CreateIterator(); + EXPECT_EQ(iter->Next().value()->ts, 100); + EXPECT_EQ(iter->Next().value()->ts, 200); + EXPECT_EQ(iter->Next().value()->ts, 300); + + // Validate second chunk content + auto chunk2 = CreateTSChunkFromData(result[1]); + EXPECT_EQ(chunk2->GetCount(), 3); + iter = chunk2->CreateIterator(); + EXPECT_EQ(iter->Next().value()->ts, 400); + EXPECT_EQ(iter->Next().value()->ts, 500); + EXPECT_EQ(iter->Next().value()->ts, 600); + + // Validate third chunk content + auto chunk3 = CreateTSChunkFromData(result[2]); + EXPECT_EQ(chunk3->GetCount(), 1); + iter = chunk3->CreateIterator(); + EXPECT_EQ(iter->Next().value()->ts, 700); + + // Test case 3: Equal split mode (7 samples into 2 chunks of 4 and 3) + test_chunk = CreateTSChunkFromData(merged_data); + new_samples = {MakeSample(400, 4.0), MakeSample(500, 5.0), MakeSample(600, 6.0), MakeSample(700, 7.0)}; + new_batch = SampleBatch(new_samples, DuplicatePolicy::LAST); + result = test_chunk->UpsertSampleAndSplit(new_batch.AsSlice(), 3, false); + ASSERT_EQ(result.size(), 2); + EXPECT_EQ(result[0].size(), TSChunk::MetaData::kEncodedSize + 4 * sizeof(TSSample)); + EXPECT_EQ(result[1].size(), TSChunk::MetaData::kEncodedSize + 3 * sizeof(TSSample)); + + // Validate split distribution + chunk1 = CreateTSChunkFromData(result[0]); + chunk2 = CreateTSChunkFromData(result[1]); + EXPECT_EQ(chunk1->GetCount(), 4); + EXPECT_EQ(chunk2->GetCount(), 3); + EXPECT_EQ(chunk1->GetFirstTimestamp(), 100); + EXPECT_EQ(chunk1->GetLastTimestamp(), 400); + EXPECT_EQ(chunk2->GetFirstTimestamp(), 500); + EXPECT_EQ(chunk2->GetLastTimestamp(), 700); + + // Test case 4: Equal split mode (no split) + test_chunk = CreateTSChunkFromData(merged_data); + new_samples = {MakeSample(400, 4.0), MakeSample(500, 5.0), MakeSample(600, 6.0), MakeSample(700, 7.0)}; + new_batch = SampleBatch(new_samples, DuplicatePolicy::LAST); + result = test_chunk->UpsertSampleAndSplit(new_batch.AsSlice(), 4, false); + EXPECT_EQ(result.size(), 1); +} + } // namespace test diff --git a/tests/cppunit/types/timeseries_test.cc b/tests/cppunit/types/timeseries_test.cc new file mode 100644 index 00000000000..5366bff7eed --- /dev/null +++ b/tests/cppunit/types/timeseries_test.cc @@ -0,0 +1,106 @@ +/* + * 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 + +#include +#include + +#include "test_base.h" +#include "types/redis_timeseries.h" + +class TimeSeriesTest : public TestBase { + protected: + explicit TimeSeriesTest() = default; + ~TimeSeriesTest() override = default; + + void SetUp() override { + key_ = "test_ts_key"; + ts_db_ = std::make_unique(storage_.get(), "ts_namespace"); + } + + std::unique_ptr ts_db_; +}; + +TEST_F(TimeSeriesTest, Create) { + redis::TSCreateOption option; + option.retention_time = 3600; + option.chunk_size = 1024; + option.chunk_type = TimeSeriesMetadata::ChunkType::COMPRESSED; + + auto s = ts_db_->Create(*ctx_, key_, option); + EXPECT_TRUE(s.ok()); + + s = ts_db_->Create(*ctx_, key_, option); + EXPECT_FALSE(s.ok()); + EXPECT_EQ(s.ToString(), "Invalid argument: key already exists"); +} + +TEST_F(TimeSeriesTest, Add) { + redis::TSCreateOption option; + option.retention_time = 3600; + option.chunk_size = 3; + auto s = ts_db_->Create(*ctx_, key_, option); + EXPECT_TRUE(s.ok()); + + TSSample sample{1620000000, 123.45}; + TSChunk::AddResultWithTS result; + s = ts_db_->Add(*ctx_, key_, sample, option, &result); + EXPECT_TRUE(s.ok()); + EXPECT_EQ(result.first, TSChunk::AddResult::kOk); + EXPECT_EQ(result.second, sample.ts); +} + +TEST_F(TimeSeriesTest, MAdd) { + redis::TSCreateOption option; + option.retention_time = 10; + auto s = ts_db_->Create(*ctx_, key_, option); + EXPECT_TRUE(s.ok()); + + std::vector samples = {{1, 10}, {3, 10}, {2, 20}, {3, 20}, {4, 20}, {13, 20}, {1, 20}, {14, 20}}; + std::vector results; + results.resize(samples.size()); + + s = ts_db_->MAdd(*ctx_, key_, samples, &results); + EXPECT_TRUE(s.ok()); + + // Expected results: kOk/kBlock/kOld verification + std::vector expected_results = {TSChunk::AddResult::kOk, // 1 + TSChunk::AddResult::kOk, // 3 + TSChunk::AddResult::kOk, // 2 + TSChunk::AddResult::kBlock, // duplicate 3 + TSChunk::AddResult::kOk, // 4 + TSChunk::AddResult::kOk, // 13 + TSChunk::AddResult::kOld, // 1 (older than retention) + TSChunk::AddResult::kOk}; // 14 + + std::vector expected_ts = {1, 3, 2, 0, 4, 13, 0, 14}; + + for (size_t i = 0; i < results.size(); ++i) { + EXPECT_EQ(results[i].first, expected_results[i]) << "Result mismatch at index " << i; + if (expected_results[i] == TSChunk::AddResult::kOk) { + EXPECT_EQ(results[i].second, expected_ts[i]) << "Timestamp mismatch at index " << i; + } + } + s = ts_db_->MAdd(*ctx_, key_, {{14, 0}}, &results); + EXPECT_TRUE(s.ok()); + EXPECT_EQ(results.size(), 1); + EXPECT_EQ(results[0].first, TSChunk::AddResult::kBlock); +} diff --git a/tests/gocase/unit/type/timeseries/timeseries_test.go b/tests/gocase/unit/type/timeseries/timeseries_test.go new file mode 100644 index 00000000000..9d4f1b422d6 --- /dev/null +++ b/tests/gocase/unit/type/timeseries/timeseries_test.go @@ -0,0 +1,147 @@ +/* + * 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. + */ +package timeseries + +import ( + "context" + "strconv" + "testing" + "time" + + "github.com/apache/kvrocks/tests/gocase/util" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestTimeSeries(t *testing.T) { + configOptions := []util.ConfigOptions{ + { + Name: "txn-context-enabled", + Options: []string{"yes", "no"}, + ConfigType: util.YesNo, + }, + } + + configsMatrix, err := util.GenerateConfigsMatrix(configOptions) + require.NoError(t, err) + + for _, configs := range configsMatrix { + testTimeSeries(t, configs) + } +} + +func testTimeSeries(t *testing.T, configs util.KvrocksServerConfigs) { + srv := util.StartServer(t, configs) + defer srv.Close() + ctx := context.Background() + rdb := srv.NewClient() + defer func() { require.NoError(t, rdb.Close()) }() + + key := "test_ts_key" + t.Run("TS.CREATE Basic Creation", func(t *testing.T) { + require.NoError(t, rdb.Del(ctx, key).Err()) + require.NoError(t, rdb.Do(ctx, "ts.create", key, "retention", "3600", "chunk_size", "2048", "encoding", "uncompressed", "duplicate_policy", "last", "labels", "label1", "value1").Err()) + }) + + t.Run("TS.CREATE Invalid RETENTION", func(t *testing.T) { + require.NoError(t, rdb.Del(ctx, key).Err()) + require.ErrorContains(t, rdb.Do(ctx, "ts.create", key, "retention", "abc").Err(), "Couldn't parse RETENTION") + require.ErrorContains(t, rdb.Do(ctx, "ts.create", key, "retention", "-100").Err(), "Couldn't parse RETENTION") + }) + + t.Run("TS.CREATE Invalid CHUNK_SIZE", func(t *testing.T) { + require.NoError(t, rdb.Del(ctx, key).Err()) + require.ErrorContains(t, rdb.Do(ctx, "ts.create", key, "chunk_size", "abc").Err(), "invalid CHUNK_SIZE") + require.ErrorContains(t, rdb.Do(ctx, "ts.create", key, "chunk_size", "-1024").Err(), "invalid CHUNK_SIZE") + }) + + t.Run("TS.CREATE Invalid ENCODING", func(t *testing.T) { + require.NoError(t, rdb.Del(ctx, key).Err()) + require.ErrorContains(t, rdb.Do(ctx, "ts.create", key, "encoding", "invalid").Err(), "unknown ENCODING parameter") + }) + + t.Run("TS.CREATE Invalid DUPLICATE_POLICY", func(t *testing.T) { + require.NoError(t, rdb.Del(ctx, key).Err()) + require.ErrorContains(t, rdb.Do(ctx, "ts.create", key, "duplicate_policy", "invalid").Err(), "Unknown DUPLICATE_POLICY") + }) + + t.Run("TS.ADD Basic Add", func(t *testing.T) { + require.NoError(t, rdb.Del(ctx, key).Err()) + require.NoError(t, rdb.Do(ctx, "ts.create", key).Err()) + require.Equal(t, int64(1000), rdb.Do(ctx, "ts.add", key, "1000", "12.3").Val()) + require.Equal(t, int64(1000), rdb.Do(ctx, "ts.add", "autocreate", "1000", "12.3").Val()) + }) + + t.Run("TS.ADD Invalid Timestamp", func(t *testing.T) { + require.NoError(t, rdb.Del(ctx, key).Err()) + require.ErrorContains(t, rdb.Do(ctx, "ts.add", key, "abc", "12.3").Err(), "invalid timestamp") + require.ErrorContains(t, rdb.Do(ctx, "ts.add", key, "-100", "12.3").Err(), "invalid timestamp") + }) + + t.Run("TS.ADD Invalid Value", func(t *testing.T) { + require.NoError(t, rdb.Del(ctx, key).Err()) + require.ErrorContains(t, rdb.Do(ctx, "ts.add", key, "1000", "abc").Err(), "invalid value") + }) + + t.Run("TS.ADD Duplicate Policy Block", func(t *testing.T) { + require.NoError(t, rdb.Del(ctx, key).Err()) + require.NoError(t, rdb.Do(ctx, "ts.create", key, "duplicate_policy", "block").Err()) + require.Equal(t, int64(1000), rdb.Do(ctx, "ts.add", key, "1000", "12.3").Val()) + require.ErrorContains(t, rdb.Do(ctx, "ts.add", key, "1000", "13.4").Err(), "update is not supported when DUPLICATE_POLICY is set to BLOCK mode") + }) + + t.Run("TS.ADD With Retention", func(t *testing.T) { + require.NoError(t, rdb.Del(ctx, key).Err()) + require.NoError(t, rdb.Do(ctx, "ts.create", key, "retention", "1000").Err()) + currentTs := time.Now().UnixMilli() + require.Equal(t, int64(currentTs), rdb.Do(ctx, "ts.add", key, strconv.FormatInt(currentTs, 10), "12.3").Val()) + oldTs := currentTs - 2000 + require.ErrorContains(t, rdb.Do(ctx, "ts.add", key, strconv.FormatInt(oldTs, 10), "12.3").Err(), "Timestamp is older than retention") + }) + + t.Run("TS.MADD Basic Test", func(t *testing.T) { + require.NoError(t, rdb.Del(ctx, key).Err()) + require.NoError(t, rdb.Do(ctx, "ts.create", key).Err()) + require.Equal(t, []interface{}{int64(1000), int64(2000)}, rdb.Do(ctx, "ts.madd", key, "1000", "12.3", key, "2000", "13.4").Val()) + }) + + t.Run("TS.MADD Invalid Arguments", func(t *testing.T) { + require.NoError(t, rdb.Del(ctx, key).Err()) + require.ErrorContains(t, rdb.Do(ctx, "ts.madd", key, "abc", "12.3").Err(), "invalid timestamp") + require.ErrorContains(t, rdb.Do(ctx, "ts.madd", key, "1000", "12.3", "invalidkey").Err(), "wrong number of arguments") + }) + + t.Run("TS.MADD Duplicate Handling", func(t *testing.T) { + require.NoError(t, rdb.Del(ctx, key).Err()) + require.NoError(t, rdb.Do(ctx, "ts.create", key, "duplicate_policy", "block").Err()) + require.Equal(t, int64(1000), rdb.Do(ctx, "ts.add", key, "1000", "12.3").Val()) + res := rdb.Do(ctx, "ts.madd", key, "1000", "13.4", key, "1000", "14.5").Val().([]interface{}) + assert.Contains(t, res[0], "update is not supported when DUPLICATE_POLICY is set to BLOCK mode") + assert.Contains(t, res[1], "update is not supported when DUPLICATE_POLICY is set to BLOCK mode") + }) + + t.Run("TS.MADD Nonexistent Key", func(t *testing.T) { + require.NoError(t, rdb.Del(ctx, "nonexistent").Err()) + require.NoError(t, rdb.Del(ctx, "existent").Err()) + require.NoError(t, rdb.Do(ctx, "ts.create", "existent").Err()) + res := rdb.Do(ctx, "ts.madd", "nonexistent", "1000", "12.3", "existent", "1000", "13.4").Val().([]interface{}) + assert.Contains(t, res[0], "the key is not a TSDB key") + assert.Equal(t, res[1], int64(1000)) + }) +} From 9d6c532133f42bd92bc31d0fd7272c63157a346c Mon Sep 17 00:00:00 2001 From: RX Xiao <74295100+yezhizi@users.noreply.github.com> Date: Fri, 22 Aug 2025 10:52:15 +0800 Subject: [PATCH 17/43] feat(ts): Add `TS.INFO` command (#3133) --- src/commands/cmd_timeseries.cc | 108 +++++++++++++++++- src/types/redis_timeseries.cc | 101 +++++++++++++++- src/types/redis_timeseries.h | 14 +++ .../unit/type/timeseries/timeseries_test.go | 84 ++++++++++++++ 4 files changed, 300 insertions(+), 7 deletions(-) diff --git a/src/commands/cmd_timeseries.cc b/src/commands/cmd_timeseries.cc index b993d3b0c5a..368fbb04e77 100644 --- a/src/commands/cmd_timeseries.cc +++ b/src/commands/cmd_timeseries.cc @@ -36,6 +36,25 @@ constexpr const char *errDupBlock = "Error at upsert, update is not supported when DUPLICATE_POLICY is set to BLOCK mode"; constexpr const char *errTSKeyNotFound = "the key is not a TSDB key"; +using ChunkType = TimeSeriesMetadata::ChunkType; +using DuplicatePolicy = TimeSeriesMetadata::DuplicatePolicy; +using TSAggregatorType = redis::TSAggregatorType; + +const std::unordered_map kChunkTypeMap = { + {ChunkType::COMPRESSED, "compressed"}, + {ChunkType::UNCOMPRESSED, "uncompressed"}, +}; +const std::unordered_map kDuplicatePolicyMap = { + {DuplicatePolicy::BLOCK, "block"}, {DuplicatePolicy::FIRST, "first"}, {DuplicatePolicy::LAST, "last"}, + {DuplicatePolicy::MIN, "min"}, {DuplicatePolicy::MAX, "max"}, {DuplicatePolicy::SUM, "sum"}, +}; +const std::unordered_map kAggregatorTypeMap = { + {TSAggregatorType::AVG, "avg"}, {TSAggregatorType::SUM, "sum"}, {TSAggregatorType::MIN, "min"}, + {TSAggregatorType::MAX, "max"}, {TSAggregatorType::RANGE, "range"}, {TSAggregatorType::COUNT, "count"}, + {TSAggregatorType::FIRST, "first"}, {TSAggregatorType::LAST, "last"}, {TSAggregatorType::STD_P, "std.p"}, + {TSAggregatorType::STD_S, "std.s"}, {TSAggregatorType::VAR_P, "var.p"}, {TSAggregatorType::VAR_S, "var.s"}, +}; + std::string FormatAddResultAsRedisReply(TSChunk::AddResultWithTS res) { using AddResult = TSChunk::AddResult; switch (res.first) { @@ -51,6 +70,30 @@ std::string FormatAddResultAsRedisReply(TSChunk::AddResultWithTS res) { return ""; } +std::string_view FormatChunkTypeAsRedisReply(ChunkType chunk_type) { + auto it = kChunkTypeMap.find(chunk_type); + if (it == kChunkTypeMap.end()) { + unreachable(); + } + return it->second; +} + +std::string_view FormatDuplicatePolicyAsRedisReply(DuplicatePolicy policy) { + auto it = kDuplicatePolicyMap.find(policy); + if (it == kDuplicatePolicyMap.end()) { + unreachable(); + } + return it->second; +} + +std::string_view FormatAggregatorTypeAsRedisReply(TSAggregatorType aggregator) { + auto it = kAggregatorTypeMap.find(aggregator); + if (it == kAggregatorTypeMap.end()) { + unreachable(); + } + return it->second; +} + } // namespace namespace redis { @@ -112,8 +155,6 @@ class CommandTSCreateBase : public KeywordCommandBase { } protected: - using DuplicatePolicy = TimeSeriesMetadata::DuplicatePolicy; - const TSCreateOption &getCreateOption() const { return create_option_; } private: @@ -138,7 +179,6 @@ class CommandTSCreateBase : public KeywordCommandBase { } Status handleEncoding(TSOptionsParser &parser) { - using ChunkType = TimeSeriesMetadata::ChunkType; if (parser.EatEqICase("UNCOMPRESSED")) { create_option_.chunk_type = ChunkType::UNCOMPRESSED; } else if (parser.EatEqICase("COMPRESSED")) { @@ -200,6 +240,65 @@ class CommandTSCreate : public CommandTSCreateBase { } }; +class CommandTSInfo : public Commander { + public: + Status Parse(const std::vector &args) override { + user_key_ = args[1]; + return Commander::Parse(args); + } + Status Execute(engine::Context &ctx, Server *srv, Connection *conn, std::string *output) override { + auto timeseries_db = TimeSeries(srv->storage, conn->GetNamespace()); + TSInfoResult info; + auto s = timeseries_db.Info(ctx, user_key_, &info); + if (s.IsNotFound()) return {Status::RedisExecErr, errTSKeyNotFound}; + if (!s.ok()) return {Status::RedisExecErr, s.ToString()}; + *output = redis::MultiLen(24); + *output += redis::SimpleString("totalSamples"); + *output += redis::Integer(info.total_samples); + *output += redis::SimpleString("memoryUsage"); + *output += redis::Integer(info.memory_usage); + *output += redis::SimpleString("firstTimestamp"); + *output += redis::Integer(info.first_timestamp); + *output += redis::SimpleString("lastTimestamp"); + *output += redis::Integer(info.last_timestamp); + *output += redis::SimpleString("retentionTime"); + *output += redis::Integer(info.metadata.retention_time); + *output += redis::SimpleString("chunkCount"); + *output += redis::Integer(info.metadata.size); + *output += redis::SimpleString("chunkSize"); + *output += redis::Integer(info.metadata.chunk_size); + *output += redis::SimpleString("chunkType"); + *output += redis::SimpleString(FormatChunkTypeAsRedisReply(info.metadata.chunk_type)); + *output += redis::SimpleString("duplicatePolicy"); + *output += redis::SimpleString(FormatDuplicatePolicyAsRedisReply(info.metadata.duplicate_policy)); + *output += redis::SimpleString("labels"); + std::vector labels_str; + labels_str.reserve(info.labels.size()); + for (const auto &label : info.labels) { + auto str = redis::Array({redis::BulkString(label.k), redis::BulkString(label.v)}); + labels_str.push_back(str); + } + *output += redis::Array(labels_str); + *output += redis::SimpleString("sourceKey"); + *output += info.metadata.source_key.empty() ? redis::NilString(redis::RESP::v3) + : redis::BulkString(info.metadata.source_key); + *output += redis::SimpleString("rules"); + std::vector rules_str; + rules_str.reserve(info.downstream_rules.size()); + for (const auto &rule : info.downstream_rules) { + auto str = redis::Array({redis::BulkString(rule.first), redis::Integer(rule.second.bucket_duration), + redis::SimpleString(FormatAggregatorTypeAsRedisReply(rule.second.aggregator)), + redis::Integer(rule.second.alignment)}); + rules_str.push_back(str); + } + *output += redis::Array(rules_str); + return Status::OK(); + } + + private: + std::string user_key_; +}; + class CommandTSAdd : public CommandTSCreateBase { public: CommandTSAdd() : CommandTSCreateBase(4, 0) { @@ -324,6 +423,7 @@ class CommandTSMAdd : public Commander { REDIS_REGISTER_COMMANDS(Timeseries, MakeCmdAttr("ts.create", -2, "write", 1, 1, 1), MakeCmdAttr("ts.add", -4, "write", 1, 1, 1), - MakeCmdAttr("ts.madd", -4, "write", 1, -3, 1), ); + MakeCmdAttr("ts.madd", -4, "write", 1, -3, 1), + MakeCmdAttr("ts.info", -2, "read-only", 1, 1, 1)); } // namespace redis diff --git a/src/types/redis_timeseries.cc b/src/types/redis_timeseries.cc index 55c299a78f6..eb0bdb128ca 100644 --- a/src/types/redis_timeseries.cc +++ b/src/types/redis_timeseries.cc @@ -132,7 +132,7 @@ rocksdb::Status TimeSeries::createTimeSeries(engine::Context &ctx, const Slice & s = batch->Put(metadata_cf_handle_, ns_key, bytes); if (!s.ok()) return s; - if (!option && !option->labels.empty()) { + if (option && !option->labels.empty()) { createLabelIndexInBatch(ns_key, *metadata_out, batch, option->labels); } @@ -234,7 +234,9 @@ rocksdb::Status TimeSeries::upsertCommon(engine::Context &ctx, const Slice &ns_k s = batch->Put(new_key, new_data); if (!s.ok()) return s; } - chunk_count += new_data_list.size() - 1; + if (!new_data_list.empty()) { + chunk_count += new_data_list.size() - 1; + } } // Process samples added to latest chunk(unseal) @@ -252,7 +254,9 @@ rocksdb::Status TimeSeries::upsertCommon(engine::Context &ctx, const Slice &ns_k s = batch->Put(new_key, new_data); if (!s.ok()) return s; } - chunk_count += new_data_list.size() - (metadata.size == 0 ? 0 : 1); + if (!new_data_list.empty()) { + chunk_count += new_data_list.size() - (metadata.size == 0 ? 0 : 1); + } if (chunk_count != metadata.size) { metadata.size = chunk_count; std::string bytes; @@ -275,6 +279,26 @@ rocksdb::Status TimeSeries::createLabelIndexInBatch(const Slice &ns_key, const T return rocksdb::Status::OK(); } +rocksdb::Status TimeSeries::getLabelKVList(engine::Context &ctx, const Slice &ns_key, + const TimeSeriesMetadata &metadata, LabelKVList *labels) { + // In the emun `TSSubkeyType`, `DOWNSTREAM` is the next of `LABEL` + std::string label_upper_bound = internalKeyFromDownstreamKey(ns_key, metadata, ""); + std::string prefix = internalKeyFromLabelKey(ns_key, metadata, ""); + + rocksdb::ReadOptions read_options = ctx.DefaultScanOptions(); + rocksdb::Slice upper_bound(label_upper_bound); + read_options.iterate_upper_bound = &upper_bound; + rocksdb::Slice lower_bound(prefix); + read_options.iterate_lower_bound = &lower_bound; + + auto iter = util::UniqueIterator(ctx, read_options); + labels->clear(); + for (iter->Seek(lower_bound); iter->Valid(); iter->Next()) { + labels->push_back({labelKeyFromInternalKey(iter->key()), iter->value().ToString()}); + } + return rocksdb::Status::OK(); +} + std::string TimeSeries::internalKeyFromChunkID(const Slice &ns_key, const TimeSeriesMetadata &metadata, uint64_t id) const { std::string sub_key; @@ -312,6 +336,13 @@ uint64_t TimeSeries::chunkIDFromInternalKey(Slice internal_key) { return DecodeFixed64(internal_key.data()); } +std::string TimeSeries::labelKeyFromInternalKey(Slice internal_key) const { + auto key = InternalKey(internal_key, storage_->IsSlotIdEncoded()); + auto label_key = key.GetSubKey(); + label_key.remove_prefix(sizeof(TSSubkeyType)); + return label_key.ToString(); +} + rocksdb::Status TimeSeries::Create(engine::Context &ctx, const Slice &user_key, const TSCreateOption &option) { std::string ns_key = AppendNamespacePrefix(user_key); @@ -359,4 +390,68 @@ rocksdb::Status TimeSeries::MAdd(engine::Context &ctx, const Slice &user_key, st return rocksdb::Status::OK(); } +rocksdb::Status TimeSeries::Info(engine::Context &ctx, const Slice &user_key, TSInfoResult *res) { + std::string ns_key = AppendNamespacePrefix(user_key); + + rocksdb::Status s = getTimeSeriesMetadata(ctx, ns_key, &res->metadata); + if (!s.ok()) { + return s; + } + auto chunk_count = res->metadata.size; + auto &metadata = res->metadata; + // Approximate total samples + res->total_samples = chunk_count * res->metadata.chunk_size; + // TODO: Estimate disk usage for the field `memoryUsage` + res->memory_usage = 0; + // Retrieve the first and last timestamp + std::string chunk_upper_bound = internalKeyFromLabelKey(ns_key, metadata, ""); + std::string end_key = internalKeyFromChunkID(ns_key, metadata, TSSample::MAX_TIMESTAMP); + std::string prefix = end_key.substr(0, end_key.size() - sizeof(uint64_t)); + + rocksdb::ReadOptions read_options = ctx.DefaultScanOptions(); + rocksdb::Slice upper_bound(chunk_upper_bound); + read_options.iterate_upper_bound = &upper_bound; + rocksdb::Slice lower_bound(prefix); + read_options.iterate_lower_bound = &lower_bound; + + auto iter = util::UniqueIterator(ctx, read_options); + iter->SeekForPrev(end_key); + if (!iter->Valid() || !iter->key().starts_with(prefix)) { + // no chunk + res->first_timestamp = 0; + res->last_timestamp = 0; + } else { + auto chunk = CreateTSChunkFromData(iter->value()); + res->last_timestamp = chunk->GetLastTimestamp(); + uint64_t retention_bound = (metadata.retention_time > 0 && res->last_timestamp > metadata.retention_time) + ? res->last_timestamp - metadata.retention_time + : 0; + auto bound_key = internalKeyFromChunkID(ns_key, metadata, retention_bound); + iter->SeekForPrev(bound_key); + if (!iter->Valid() || !iter->key().starts_with(prefix)) { + if (!iter->Valid()) { + iter->Seek(bound_key); + } else { + iter->Next(); + } + chunk = CreateTSChunkFromData(iter->value()); + res->first_timestamp = chunk->GetFirstTimestamp(); + } else { + chunk = CreateTSChunkFromData(iter->value()); + auto chunk_it = chunk->CreateIterator(); + while (chunk_it->HasNext()) { + auto sample = chunk_it->Next().value(); + if (sample->ts >= retention_bound) { + res->first_timestamp = sample->ts; + break; + } + } + } + } + getLabelKVList(ctx, ns_key, metadata, &res->labels); + // TODO: Retrieve downstream downstream_rules + + return rocksdb::Status::OK(); +} + } // namespace redis diff --git a/src/types/redis_timeseries.h b/src/types/redis_timeseries.h index c853ec83363..16acdacf71c 100644 --- a/src/types/redis_timeseries.h +++ b/src/types/redis_timeseries.h @@ -106,6 +106,16 @@ struct TSCreateOption { TSCreateOption(); }; +struct TSInfoResult { + TimeSeriesMetadata metadata; + uint64_t total_samples; + uint64_t memory_usage; + uint64_t first_timestamp; + uint64_t last_timestamp; + std::vector> downstream_rules; + LabelKVList labels; +}; + TimeSeriesMetadata CreateMetadataFromOption(const TSCreateOption &option); class TimeSeries : public SubKeyScanner { @@ -120,6 +130,7 @@ class TimeSeries : public SubKeyScanner { AddResultWithTS *res, const DuplicatePolicy *on_dup_policy = nullptr); rocksdb::Status MAdd(engine::Context &ctx, const Slice &user_key, std::vector samples, std::vector *res); + rocksdb::Status Info(engine::Context &ctx, const Slice &user_key, TSInfoResult *res); private: rocksdb::Status getTimeSeriesMetadata(engine::Context &ctx, const Slice &ns_key, TimeSeriesMetadata *metadata); @@ -127,6 +138,8 @@ class TimeSeries : public SubKeyScanner { const TSCreateOption *options); rocksdb::Status getOrCreateTimeSeries(engine::Context &ctx, const Slice &ns_key, TimeSeriesMetadata *metadata_out, const TSCreateOption *option = nullptr); + rocksdb::Status getLabelKVList(engine::Context &ctx, const Slice &ns_key, const TimeSeriesMetadata &metadata, + LabelKVList *labels); rocksdb::Status upsertCommon(engine::Context &ctx, const Slice &ns_key, TimeSeriesMetadata &metadata, SampleBatch &sample_batch); rocksdb::Status createLabelIndexInBatch(const Slice &ns_key, const TimeSeriesMetadata &metadata, @@ -136,6 +149,7 @@ class TimeSeries : public SubKeyScanner { std::string internalKeyFromLabelKey(const Slice &ns_key, const TimeSeriesMetadata &metadata, Slice label_key) const; std::string internalKeyFromDownstreamKey(const Slice &ns_key, const TimeSeriesMetadata &metadata, Slice downstream_key) const; + std::string labelKeyFromInternalKey(Slice internal_key) const; static uint64_t chunkIDFromInternalKey(Slice internal_key); }; diff --git a/tests/gocase/unit/type/timeseries/timeseries_test.go b/tests/gocase/unit/type/timeseries/timeseries_test.go index 9d4f1b422d6..ed23a3dfc5f 100644 --- a/tests/gocase/unit/type/timeseries/timeseries_test.go +++ b/tests/gocase/unit/type/timeseries/timeseries_test.go @@ -81,6 +81,90 @@ func testTimeSeries(t *testing.T, configs util.KvrocksServerConfigs) { require.ErrorContains(t, rdb.Do(ctx, "ts.create", key, "duplicate_policy", "invalid").Err(), "Unknown DUPLICATE_POLICY") }) + // Test non-existent key + t.Run("TS.INFO Non-Existent Key", func(t *testing.T) { + _, err := rdb.Do(ctx, "ts.info", "test_info_key").Result() + require.ErrorContains(t, err, "the key is not a TSDB key") + }) + + t.Run("TS.INFO Initial State", func(t *testing.T) { + key := "test_info_key" + // Create timeseries with custom options + require.NoError(t, rdb.Do(ctx, "ts.create", key, "retention", "10", "chunk_size", "3", + "labels", "k1", "v1", "k2", "v2").Err()) + vals, err := rdb.Do(ctx, "ts.info", key).Slice() + require.NoError(t, err) + require.Equal(t, 24, len(vals)) + + // totalSamples = 0 + require.Equal(t, "totalSamples", vals[0]) + require.Equal(t, int64(0), vals[1]) + + // memoryUsage = 0 + require.Equal(t, "memoryUsage", vals[2]) + require.Equal(t, int64(0), vals[3]) + + // retentionTime = 10 + require.Equal(t, "retentionTime", vals[8]) + require.Equal(t, int64(10), vals[9]) + + // chunkSize = 3 + require.Equal(t, "chunkSize", vals[12]) + require.Equal(t, int64(3), vals[13]) + + // chunkType = uncompressed + require.Equal(t, "chunkType", vals[14]) + require.Equal(t, "uncompressed", vals[15]) + + // duplicatePolicy = block + require.Equal(t, "duplicatePolicy", vals[16]) + require.Equal(t, "block", vals[17]) + + // labels = [(k1,v1), (k2,v2)] + require.Equal(t, "labels", vals[18]) + labels := vals[19].([]interface{}) + require.Equal(t, 2, len(labels)) + for i, expected := range [][]string{{"k1", "v1"}, {"k2", "v2"}} { + pair := labels[i].([]interface{}) + require.Equal(t, expected[0], pair[0]) + require.Equal(t, expected[1], pair[1]) + } + + // sourceKey = nil + require.Equal(t, "sourceKey", vals[20]) + require.Nil(t, []byte(nil), vals[21]) + + // rules = empty array + require.Equal(t, "rules", vals[22]) + require.Empty(t, vals[23]) + }) + + t.Run("TS.INFO After Adding Data", func(t *testing.T) { + key := "test_info_key" + // Add samples + require.NoError(t, rdb.Do(ctx, "ts.madd", key, "1", "10", key, "3", "10", key, "2", "20", + key, "3", "20", key, "4", "20", key, "13", "20", key, "1", "20", key, "14", "20").Err()) + + vals, err := rdb.Do(ctx, "ts.info", key).Slice() + require.NoError(t, err) + + // totalSamples = 6 + require.Equal(t, "totalSamples", vals[0]) + require.Equal(t, int64(6), vals[1]) + + // firstTimestamp = 4 (earliest after retention) + require.Equal(t, "firstTimestamp", vals[4]) + require.Equal(t, int64(4), vals[5]) + + // lastTimestamp = 14 + require.Equal(t, "lastTimestamp", vals[6]) + require.Equal(t, int64(14), vals[7]) + + // chunkCount = 2 + require.Equal(t, "chunkCount", vals[10]) + require.Equal(t, int64(2), vals[11]) + }) + t.Run("TS.ADD Basic Add", func(t *testing.T) { require.NoError(t, rdb.Del(ctx, key).Err()) require.NoError(t, rdb.Do(ctx, "ts.create", key).Err()) From ff658f805597779dd995e1bdd7aa2642d18d473b Mon Sep 17 00:00:00 2001 From: Twice Date: Fri, 22 Aug 2025 11:00:43 +0800 Subject: [PATCH 18/43] chore(.asf.yaml): enable auto merge and disable wiki (#3137) --- .asf.yaml | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/.asf.yaml b/.asf.yaml index 9034a87aa1c..436fb11926f 100644 --- a/.asf.yaml +++ b/.asf.yaml @@ -29,10 +29,19 @@ github: - namespace - redis - redis-cluster + pull_requests: + allow_auto_merge: true + allow_update_branch: true enabled_merge_buttons: squash: true + squash_commit_message: PR_TITLE_AND_DESC merge: false rebase: true + features: + wiki: false + issues: true + projects: true + discussions: true protected_branches: unstable: required_pull_request_reviews: From 53e82f81652f7ec215d12997a4b91273dead0f97 Mon Sep 17 00:00:00 2001 From: Jonah Gao Date: Fri, 22 Aug 2025 12:33:02 +0800 Subject: [PATCH 19/43] chore: remove unused `autoResizeBlockAndSST` method and config (#3136) The usage of this function was removed in https://github.com/apache/kvrocks/pull/395. Co-authored-by: Twice --- src/config/config.cc | 1 - src/config/config.h | 1 - src/server/server.cc | 85 -------------------------------------------- src/server/server.h | 1 - 4 files changed, 88 deletions(-) diff --git a/src/config/config.cc b/src/config/config.cc index b8b54b32a37..f17db782925 100644 --- a/src/config/config.cc +++ b/src/config/config.cc @@ -219,7 +219,6 @@ Config::Config() { spdlog::level::off)}, {"purge-backup-on-fullsync", false, new YesNoField(&purge_backup_on_fullsync, false)}, {"rename-command", true, new MultiStringField(&rename_command_, std::vector{})}, - {"auto-resize-block-and-sst", false, new YesNoField(&auto_resize_block_and_sst, true)}, {"fullsync-recv-file-delay", false, new IntField(&fullsync_recv_file_delay, 0, 0, INT_MAX)}, {"cluster-enabled", true, new YesNoField(&cluster_enabled, false)}, {"migrate-speed", false, new IntField(&migrate_speed, 4096, 0, INT_MAX)}, diff --git a/src/config/config.h b/src/config/config.h index 871f13236ca..578843c39c4 100644 --- a/src/config/config.h +++ b/src/config/config.h @@ -129,7 +129,6 @@ struct Config { int max_bitmap_to_string_mb = 16; bool master_use_repl_port = false; bool purge_backup_on_fullsync = false; - bool auto_resize_block_and_sst = true; int fullsync_recv_file_delay = 0; bool use_rsid_psync = false; bool replication_group_sync = false; diff --git a/src/server/server.cc b/src/server/server.cc index b6e10cbefae..40fc79f7594 100644 --- a/src/server/server.cc +++ b/src/server/server.cc @@ -1584,91 +1584,6 @@ Status Server::AsyncScanDBSize(const std::string &ns) { }); } -Status Server::autoResizeBlockAndSST() { - auto total_size = storage->GetTotalSize(kDefaultNamespace); - uint64_t total_keys = 0, estimate_keys = 0; - for (const auto &cf_handle : *storage->GetCFHandles()) { - storage->GetDB()->GetIntProperty(cf_handle, rocksdb::DB::Properties::kEstimateNumKeys, &estimate_keys); - total_keys += estimate_keys; - } - - if (total_size == 0 || total_keys == 0) { - return Status::OK(); - } - - auto average_kv_size = total_size / total_keys; - int target_file_size_base = 0; - int block_size = 0; - if (average_kv_size > 512 * KiB) { - target_file_size_base = 1024; - block_size = 1 * MiB; - } else if (average_kv_size > 256 * KiB) { - target_file_size_base = 512; - block_size = 512 * KiB; - } else if (average_kv_size > 32 * KiB) { - target_file_size_base = 256; - block_size = 256 * KiB; - } else if (average_kv_size > 1 * KiB) { - target_file_size_base = 128; - block_size = 32 * KiB; - } else if (average_kv_size > 128) { - target_file_size_base = 64; - block_size = 8 * KiB; - } else { - target_file_size_base = 16; - block_size = 2 * KiB; - } - - if (target_file_size_base == config_->rocks_db.target_file_size_base && - target_file_size_base == config_->rocks_db.write_buffer_size && block_size == config_->rocks_db.block_size) { - return Status::OK(); - } - - if (target_file_size_base != config_->rocks_db.target_file_size_base) { - auto old_target_file_size_base = config_->rocks_db.target_file_size_base; - auto s = config_->Set(this, "rocksdb.target_file_size_base", std::to_string(target_file_size_base)); - info( - "[server] Resize rocksdb.target_file_size_base from {} to {}, " - "average_kv_size: {}, total_size: {}, total_keys: {}, result: {}", - old_target_file_size_base, target_file_size_base, average_kv_size, total_size, total_keys, s.Msg()); - if (!s.IsOK()) { - return s; - } - } - - if (target_file_size_base != config_->rocks_db.write_buffer_size) { - auto old_write_buffer_size = config_->rocks_db.write_buffer_size; - auto s = config_->Set(this, "rocksdb.write_buffer_size", std::to_string(target_file_size_base)); - info( - "[server] Resize rocksdb.write_buffer_size from {} to {}, " - "average_kv_size: {}, total_size: {}, " - "total_keys: {}, result: {}", - old_write_buffer_size, target_file_size_base, average_kv_size, total_size, total_keys, s.Msg()); - if (!s.IsOK()) { - return s; - } - } - - if (block_size != config_->rocks_db.block_size) { - auto s = storage->SetOptionForAllColumnFamilies("table_factory.block_size", std::to_string(block_size)); - info( - "[server] Resize rocksdb.block_size from {} to {}, " - "average_kv_size: {}, total_size: {}, " - "total_keys: {}, result: {}", - config_->rocks_db.block_size, block_size, average_kv_size, total_size, total_keys, s.Msg()); - if (!s.IsOK()) { - return s; - } - - config_->rocks_db.block_size = block_size; - } - - auto s = config_->Rewrite(namespace_.List()); - info("[server] Rewrite config, result: {}", s.Msg()); - - return Status::OK(); -} - void Server::GetLatestKeyNumStats(const std::string &ns, KeyNumStats *stats) { auto iter = db_scan_infos_.find(ns); if (iter != db_scan_infos_.end()) { diff --git a/src/server/server.h b/src/server/server.h index 37211c4f4c3..54f4c2cc4a3 100644 --- a/src/server/server.h +++ b/src/server/server.h @@ -357,7 +357,6 @@ class Server { void cron(); void recordInstantaneousMetrics(); static void updateCachedTime(); - Status autoResizeBlockAndSST(); void updateWatchedKeysFromRange(const std::vector &args, const redis::CommandKeyRange &range); void updateAllWatchedKeys(); void increaseWorkerThreads(size_t delta); From 6df33094b652988eee51c14c3ef923c3b6a05d84 Mon Sep 17 00:00:00 2001 From: Twice Date: Sat, 23 Aug 2025 13:36:09 +0800 Subject: [PATCH 20/43] feat(scripting): support strict key-accessing mode for lua scripting (#3139) Accessing undeclared keys in lua scripting may lead to unexpected behavior in the current design of kvrocks (also in redis, refer to https://redis.io/docs/latest/commands/eval/), so in this PR we add a new option `lua-strict-key-accessing` to prevent users to access undeclared keys in lua, e.g. ``` EVAL "return redis.call('set', 'a', 1)" // ERROR! EVAL "return redis.call('set', KEYS[1], 1)" 1 a // ok ``` This check is performed in both lua scripting and lua functions. --- kvrocks.conf | 13 +++++++ src/config/config.cc | 1 + src/config/config.h | 2 ++ src/storage/scripting.cc | 36 +++++++++++++++++-- src/storage/scripting.h | 5 +-- tests/gocase/unit/scripting/function_test.go | 34 ++++++++++++++++++ tests/gocase/unit/scripting/scripting_test.go | 31 ++++++++++++++++ 7 files changed, 117 insertions(+), 5 deletions(-) diff --git a/kvrocks.conf b/kvrocks.conf index 67819c27438..b5693434541 100644 --- a/kvrocks.conf +++ b/kvrocks.conf @@ -420,6 +420,19 @@ txn-context-enabled no # Default: disabled # histogram-bucket-boundaries 10,20,40,60,80,100,150,250,350,500,750,1000,1500,2000,4000,8000 +# Whether the strict key-accessing mode of lua scripting is enabled. +# +# If enabled, the lua script will abort and report errors +# if it tries to access keys that are not declared in +# the script's `KEYS` table or the function's `keys` argument. +# +# Note that accessing undeclared keys may lead to unexpected behavior, +# so this option is to ensure that scripts only access keys +# that are explicitly declared. +# +# Default: no +lua-strict-key-accessing no + ################################## TLS ################################### # By default, TLS/SSL is disabled, i.e. `tls-port` is set to 0. diff --git a/src/config/config.cc b/src/config/config.cc index f17db782925..fb0f68f6770 100644 --- a/src/config/config.cc +++ b/src/config/config.cc @@ -243,6 +243,7 @@ Config::Config() { {"txn-context-enabled", true, new YesNoField(&txn_context_enabled, false)}, {"skip-block-cache-deallocation-on-close", false, new YesNoField(&skip_block_cache_deallocation_on_close, false)}, {"histogram-bucket-boundaries", true, new StringField(&histogram_bucket_boundaries_str_, "")}, + {"lua-strict-key-accessing", false, new YesNoField(&lua_strict_key_accessing, false)}, /* rocksdb options */ {"rocksdb.compression", false, diff --git a/src/config/config.h b/src/config/config.h index 578843c39c4..30527eca6d6 100644 --- a/src/config/config.h +++ b/src/config/config.h @@ -193,6 +193,8 @@ struct Config { bool skip_block_cache_deallocation_on_close = false; + bool lua_strict_key_accessing = false; + std::vector histogram_bucket_boundaries; struct RocksDB { diff --git a/src/storage/scripting.cc b/src/storage/scripting.cc index b5f865ac2ba..357b9ae2687 100644 --- a/src/storage/scripting.cc +++ b/src/storage/scripting.cc @@ -426,6 +426,9 @@ Status FunctionCall(redis::Connection *conn, engine::Context *ctx, const std::st SaveOnRegistry(lua, REGISTRY_SCRIPT_RUN_CTX_NAME, &script_run_ctx); + // save keys on registry the to perform key touching check + SaveOnRegistry(lua, REGISTRY_KEYS_NAME, &keys); + PushArray(lua, keys); PushArray(lua, argv); if (lua_pcall(lua, 2, 1, -4)) { @@ -437,6 +440,7 @@ Status FunctionCall(redis::Connection *conn, engine::Context *ctx, const std::st lua_pop(lua, 2); } + RemoveFromRegistry(lua, REGISTRY_KEYS_NAME); RemoveFromRegistry(lua, REGISTRY_SCRIPT_RUN_CTX_NAME); /* Call the Lua garbage collector from time to time to avoid a @@ -704,6 +708,9 @@ Status EvalGenericCommand(redis::Connection *conn, engine::Context *ctx, const s SetGlobalArray(lua, "KEYS", keys); SetGlobalArray(lua, "ARGV", argv); + // save keys on registry the to perform key touching check + SaveOnRegistry(lua, REGISTRY_KEYS_NAME, &keys); + if (lua_pcall(lua, 0, 1, -2)) { auto msg = fmt::format("running script (call to {}): {}", funcname, lua_tostring(lua, -1)); *output = redis::Error({Status::NotOK, msg}); @@ -719,6 +726,7 @@ Status EvalGenericCommand(redis::Connection *conn, engine::Context *ctx, const s lua_setglobal(lua, "KEYS"); lua_pushnil(lua); lua_setglobal(lua, "ARGV"); + RemoveFromRegistry(lua, REGISTRY_KEYS_NAME); RemoveFromRegistry(lua, REGISTRY_SCRIPT_RUN_CTX_NAME); /* Call the Lua garbage collector from time to time to avoid a @@ -814,6 +822,28 @@ int RedisGenericCommand(lua_State *lua, int raise_error) { return raise_error ? RaiseError(lua) : 1; } + // if it is a write command and strict mode is enabled, + // we need to check if the input keys are all in allowed keys + if (config->lua_strict_key_accessing && !(cmd_flags & redis::kCmdReadOnly)) { + auto allowed_keys = GetFromRegistry>(lua, REGISTRY_KEYS_NAME); + + attributes->ForEachKeyRange( + [&](const std::vector &args, const redis::CommandKeyRange &range) { + range.ForEachKey( + [&](const std::string &key) { + if (std::find(allowed_keys->begin(), allowed_keys->end(), key) == allowed_keys->end()) { + PushError(lua, fmt::format("Script attempted to access key '{}' which is not in the allowed keys " + "(lua-strict-key-accessing)", + key) + .c_str()); + RaiseError(lua); + } + }, + args); + }, + args); + } + if (config->cluster_enabled) { if (script_run_ctx->flags & ScriptFlagType::kScriptNoCluster) { PushError(lua, "Can not run script on cluster, 'no-cluster' flag is set"); @@ -1582,9 +1612,9 @@ Status CreateFunction(Server *srv, const std::string &body, std::string *sha, lu libname = shebang_split_sv.substr(shebang_libname_prefix.size()); if (libname.empty() || std::any_of(libname.begin(), libname.end(), [](char v) { return !std::isalnum(v) && v != '_'; })) { - return { - Status::NotOK, - "Library names can only contain letters, numbers, or underscores(_) and must be at least one character long"}; + return {Status::NotOK, + "Library names can only contain letters, numbers, or underscores(_) and must be at least one " + "character long"}; } found_libname = true; } diff --git a/src/storage/scripting.h b/src/storage/scripting.h index c28ab668916..bf03f3e104c 100644 --- a/src/storage/scripting.h +++ b/src/storage/scripting.h @@ -39,7 +39,8 @@ inline constexpr const char REDIS_LUA_REGISTER_FUNC_FLAGS_PREFIX[] = "__redis_re inline constexpr const char REDIS_FUNCTION_LIBNAME[] = "REDIS_FUNCTION_LIBNAME"; inline constexpr const char REDIS_FUNCTION_NEEDSTORE[] = "REDIS_FUNCTION_NEEDSTORE"; inline constexpr const char REDIS_FUNCTION_LIBRARIES[] = "REDIS_FUNCTION_LIBRARIES"; -inline constexpr const char REGISTRY_SCRIPT_RUN_CTX_NAME[] = "SCRIPT_RUN_CTX"; +inline constexpr const char REGISTRY_SCRIPT_RUN_CTX_NAME[] = "__SCRIPT_RUN_CTX"; +inline constexpr const char REGISTRY_KEYS_NAME[] = "__CURRENT_KEYS"; namespace lua { @@ -165,7 +166,7 @@ template void SaveOnRegistry(lua_State *lua, const char *name, T *ptr) { lua_pushstring(lua, name); if (ptr) { - lua_pushlightuserdata(lua, ptr); + lua_pushlightuserdata(lua, (void *)ptr); } else { lua_pushnil(lua); } diff --git a/tests/gocase/unit/scripting/function_test.go b/tests/gocase/unit/scripting/function_test.go index 98c92fe1adc..5d307b48860 100644 --- a/tests/gocase/unit/scripting/function_test.go +++ b/tests/gocase/unit/scripting/function_test.go @@ -581,3 +581,37 @@ func TestFunctionScriptFlags(t *testing.T) { util.ErrorRegexp(t, r.Err(), "ERR .* Script attempted to access a non local key in a cluster node script") }) } + +func TestFunctionInStrictMode(t *testing.T) { + srv := util.StartServer(t, map[string]string{}) + defer srv.Close() + + ctx := context.Background() + rdb := srv.NewClient() + defer func() { require.NoError(t, rdb.Close()) }() + + t.Run("Accessing undeclared keys in strict mode", func(t *testing.T) { + rdb.FunctionLoad(ctx, `#!lua name=tmplib + redis.register_function('set1', function(keys, args) + return redis.call('set', keys[1], args[1]) + end) + redis.register_function('set2', function(keys, args) + return redis.call('set', args[1], args[2]) + end) + `) + + rdb.ConfigSet(ctx, "lua-strict-key-accessing", "yes") + + util.ErrorRegexp(t, rdb.Do(ctx, "FCALL", "set2", 0, "x", "1").Err(), ".*'x'.*not in the allowed keys.*") + util.ErrorRegexp(t, rdb.Do(ctx, "FCALL", "set2", 1, "y", "x", "1").Err(), ".*'x'.*not in the allowed keys.*") + + require.NoError(t, rdb.Do(ctx, "FCALL", "set2", 1, "x", "x", "1").Err()) + require.NoError(t, rdb.Do(ctx, "FCALL", "set2", 2, "x", "y", "x", "1").Err()) + require.NoError(t, rdb.Do(ctx, "FCALL", "set1", 1, "x", "1").Err()) + + rdb.ConfigSet(ctx, "lua-strict-key-accessing", "no") + + require.NoError(t, rdb.Do(ctx, "FCALL", "set2", 0, "x", "1").Err()) + require.NoError(t, rdb.Do(ctx, "FCALL", "set1", 1, "x", "1").Err()) + }) +} diff --git a/tests/gocase/unit/scripting/scripting_test.go b/tests/gocase/unit/scripting/scripting_test.go index 0c886dfd5ca..ce606a04f74 100644 --- a/tests/gocase/unit/scripting/scripting_test.go +++ b/tests/gocase/unit/scripting/scripting_test.go @@ -894,3 +894,34 @@ func TestEvalScriptFlags(t *testing.T) { }) } + +func TestEvalScriptInStrictMode(t *testing.T) { + srv := util.StartServer(t, map[string]string{}) + defer srv.Close() + + ctx := context.Background() + rdb := srv.NewClient() + defer func() { require.NoError(t, rdb.Close()) }() + + t.Run("Accessing undeclared keys in strict mode", func(t *testing.T) { + rdb.ConfigSet(ctx, "lua-strict-key-accessing", "yes") + + util.ErrorRegexp(t, rdb.Eval(ctx, "return redis.call('set', 'a', 1)", []string{}).Err(), ".*'a'.*not in the allowed keys.*") + util.ErrorRegexp(t, rdb.Eval(ctx, "return redis.call('set', ARGV[1], 1)", []string{}, "a").Err(), ".*'a'.*not in the allowed keys.*") + util.ErrorRegexp(t, rdb.Eval(ctx, "return redis.call('set', 'b', 1)", []string{"a"}).Err(), ".*'b'.*not in the allowed keys.*") + util.ErrorRegexp(t, rdb.Eval(ctx, "return redis.call('set', KEYS[1]..'b', 1)", []string{"a"}).Err(), ".*'ab'.*not in the allowed keys.*") + + require.NoError(t, rdb.Eval(ctx, "return redis.call('set', KEYS[1], 1)", []string{"a"}).Err()) + require.NoError(t, rdb.Eval(ctx, "return redis.call('set', KEYS[2], 1)", []string{"a", "b"}).Err()) + require.NoError(t, rdb.Eval(ctx, "return redis.call('set', 'a', 1)", []string{"a", "b"}).Err()) + require.NoError(t, rdb.Eval(ctx, "return redis.call('set', 'b', 1)", []string{"a", "b"}).Err()) + require.NoError(t, rdb.Eval(ctx, "return redis.call('set', ARGV[1]..ARGV[2], 1)", []string{"ab"}, "a", "b").Err()) + + // read-only commands are allowed as an exception + require.NoError(t, rdb.Eval(ctx, "return redis.call('get', 'a')", []string{}).Err()) + + rdb.ConfigSet(ctx, "lua-strict-key-accessing", "no") + + require.NoError(t, rdb.Eval(ctx, "return redis.call('set', 'a', 1)", []string{}).Err()) + }) +} From 201afed607049a27fdc68b58f8c189ef400069d3 Mon Sep 17 00:00:00 2001 From: Roman Donchenko Date: Sun, 24 Aug 2025 08:41:11 +0300 Subject: [PATCH 21/43] feat(Dockerfile): add a UID for the user in the container (#3138) Also, use hardcoded IDs when creating the user and group, to ensure they remain stable. Closes #3135. Co-authored-by: Aleks Lozovyuk --- Dockerfile | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/Dockerfile b/Dockerfile index d6353b48954..f5cdf74ce14 100644 --- a/Dockerfile +++ b/Dockerfile @@ -31,13 +31,13 @@ FROM debian:bookworm-slim RUN DEBIAN_FRONTEND=noninteractive && apt-get update && apt-get upgrade -y && apt-get -y install openssl ca-certificates redis-tools binutils && apt-get clean # Create a dedicated non-root user and group -RUN groupadd -r kvrocks && useradd -r -g kvrocks kvrocks +RUN groupadd --gid=999 -r kvrocks && useradd --uid=999 -r -g kvrocks kvrocks RUN mkdir /var/run/kvrocks /var/lib/kvrocks && \ chown -R kvrocks:kvrocks /var/run/kvrocks /var/lib/kvrocks # Switch to the non-root user -USER kvrocks +USER 999 VOLUME /var/lib/kvrocks From 0851c220581de334b045732992965b87c669431f Mon Sep 17 00:00:00 2001 From: RX Xiao <74295100+yezhizi@users.noreply.github.com> Date: Mon, 25 Aug 2025 13:17:04 +0800 Subject: [PATCH 22/43] feat(ts): Add data query support and `TS.RANGE` command (#3140) Part of #3048 --- src/commands/cmd_timeseries.cc | 243 ++++++++++++- src/types/redis_timeseries.cc | 340 +++++++++++++++++- src/types/redis_timeseries.h | 79 ++-- src/types/timeseries.h | 1 + tests/cppunit/types/timeseries_test.cc | 228 ++++++++++++ .../unit/type/timeseries/timeseries_test.go | 180 ++++++++++ 6 files changed, 1040 insertions(+), 31 deletions(-) diff --git a/src/commands/cmd_timeseries.cc b/src/commands/cmd_timeseries.cc index 368fbb04e77..995c9b8728d 100644 --- a/src/commands/cmd_timeseries.cc +++ b/src/commands/cmd_timeseries.cc @@ -35,6 +35,7 @@ constexpr const char *errOldTimestamp = "Timestamp is older than retention"; constexpr const char *errDupBlock = "Error at upsert, update is not supported when DUPLICATE_POLICY is set to BLOCK mode"; constexpr const char *errTSKeyNotFound = "the key is not a TSDB key"; +constexpr const char *errTSInvalidAlign = "unknown ALIGN parameter"; using ChunkType = TimeSeriesMetadata::ChunkType; using DuplicatePolicy = TimeSeriesMetadata::DuplicatePolicy; @@ -70,6 +71,13 @@ std::string FormatAddResultAsRedisReply(TSChunk::AddResultWithTS res) { return ""; } +std::string FormatTSSampleAsRedisReply(TSSample sample) { + std::string res = redis::MultiLen(2); + res += redis::Integer(sample.ts); + res += redis::Double(redis::RESP::v3, sample.v); + return res; +} + std::string_view FormatChunkTypeAsRedisReply(ChunkType chunk_type) { auto it = kChunkTypeMap.find(chunk_type); if (it == kChunkTypeMap.end()) { @@ -285,10 +293,11 @@ class CommandTSInfo : public Commander { *output += redis::SimpleString("rules"); std::vector rules_str; rules_str.reserve(info.downstream_rules.size()); - for (const auto &rule : info.downstream_rules) { - auto str = redis::Array({redis::BulkString(rule.first), redis::Integer(rule.second.bucket_duration), - redis::SimpleString(FormatAggregatorTypeAsRedisReply(rule.second.aggregator)), - redis::Integer(rule.second.alignment)}); + for (const auto &[key, rule] : info.downstream_rules) { + const auto &aggregator = rule.aggregator; + auto str = redis::Array({redis::BulkString(key), redis::Integer(aggregator.bucket_duration), + redis::SimpleString(FormatAggregatorTypeAsRedisReply(aggregator.type)), + redis::Integer(aggregator.alignment)}); rules_str.push_back(str); } *output += redis::Array(rules_str); @@ -421,9 +430,235 @@ class CommandTSMAdd : public Commander { std::unordered_map> userkey_indexes_map_; }; +class CommandTSRangeBase : public KeywordCommandBase { + public: + CommandTSRangeBase(size_t skip_num, size_t tail_skip_num) + : KeywordCommandBase(skip_num + 2, tail_skip_num), skip_num_(skip_num) { + registerHandler("LATEST", [this](TSOptionsParser &parser) { return handleLatest(parser); }); + registerHandler("FILTER_BY_TS", [this](TSOptionsParser &parser) { return handleFilterByTS(parser); }); + registerHandler("FILTER_BY_VALUE", [this](TSOptionsParser &parser) { return handleFilterByValue(parser); }); + registerHandler("COUNT", [this](TSOptionsParser &parser) { return handleCount(parser); }); + registerHandler("ALIGN", [this](TSOptionsParser &parser) { return handleAlign(parser); }); + registerHandler("AGGREGATION", [this](TSOptionsParser &parser) { return handleAggregation(parser); }); + registerHandler("BUCKETTIMESTAMP", [this](TSOptionsParser &parser) { return handleBucketTimestamp(parser); }); + registerHandler("EMPTY", [this](TSOptionsParser &parser) { return handleEmpty(parser); }); + } + + Status Parse(const std::vector &args) override { + TSOptionsParser parser(std::next(args.begin(), static_cast(skip_num_)), args.end()); + // Parse start timestamp + auto start_ts = parser.TakeInt(); + if (!start_ts.IsOK()) { + auto start_ts_str = parser.TakeStr(); + if (!start_ts_str.IsOK() || start_ts_str.GetValue() != "-") { + return {Status::RedisParseErr, "wrong fromTimestamp"}; + } + // "-" means use default start timestamp: 0 + } else { + is_start_explicit_set_ = true; + option_.start_ts = start_ts.GetValue(); + } + + // Parse end timestamp + auto end_ts = parser.TakeInt(); + if (!end_ts.IsOK()) { + auto end_ts_str = parser.TakeStr(); + if (!end_ts_str.IsOK() || end_ts_str.GetValue() != "+") { + return {Status::RedisParseErr, "wrong toTimestamp"}; + } + // "+" means use default end timestamp: MAX_TIMESTAMP + } else { + is_end_explicit_set_ = true; + option_.end_ts = end_ts.GetValue(); + } + + auto s = KeywordCommandBase::Parse(args); + if (!s.IsOK()) return s; + if (is_alignment_explicit_set_ && option_.aggregator.type == TSAggregatorType::NONE) { + return {Status::RedisParseErr, "ALIGN parameter can only be used with AGGREGATION"}; + } + return s; + } + + const TSRangeOption &GetRangeOption() const { return option_; } + + private: + TSRangeOption option_; + size_t skip_num_; + bool is_start_explicit_set_ = false; + bool is_end_explicit_set_ = false; + bool is_alignment_explicit_set_ = false; + + Status handleLatest([[maybe_unused]] TSOptionsParser &parser) { + option_.is_return_latest = true; + return Status::OK(); + } + + Status handleFilterByTS(TSOptionsParser &parser) { + option_.filter_by_ts.clear(); + while (parser.Good()) { + auto ts = parser.TakeInt(); + if (!ts.IsOK()) break; + option_.filter_by_ts.insert(ts.GetValue()); + } + return Status::OK(); + } + + Status handleFilterByValue(TSOptionsParser &parser) { + auto min = parser.TakeFloat(); + auto max = parser.TakeFloat(); + if (!min.IsOK() || !max.IsOK()) { + return {Status::RedisParseErr, "Invalid min or max value"}; + } + option_.filter_by_value = std::make_optional(std::make_pair(min.GetValue(), max.GetValue())); + return Status::OK(); + } + + Status handleCount(TSOptionsParser &parser) { + auto count = parser.TakeInt(); + if (!count.IsOK()) { + return {Status::RedisParseErr, "Couldn't parse COUNT"}; + } + option_.count_limit = count.GetValue(); + if (option_.count_limit == 0) { + return {Status::RedisParseErr, "Invalid COUNT value"}; + } + return Status::OK(); + } + + Status handleAlign(TSOptionsParser &parser) { + auto align = parser.TakeInt(); + if (align.IsOK()) { + is_alignment_explicit_set_ = true; + option_.aggregator.alignment = align.GetValue(); + return Status::OK(); + } + + auto align_str = parser.TakeStr(); + if (!align_str.IsOK()) { + return {Status::RedisParseErr, errTSInvalidAlign}; + } + + const auto &value = align_str.GetValue(); + if (value == "-" || value == "+") { + bool is_explicit_set = value == "-" ? is_start_explicit_set_ : is_end_explicit_set_; + auto err_msg = value == "-" ? "start alignment can only be used with explicit start timestamp" + : "end alignment can only be used with explicit end timestamp"; + + if (!is_explicit_set) { + return {Status::RedisParseErr, err_msg}; + } + + option_.aggregator.alignment = value == "-" ? option_.start_ts : option_.end_ts; + } else { + return {Status::RedisParseErr, errTSInvalidAlign}; + } + is_alignment_explicit_set_ = true; + return Status::OK(); + } + + Status handleAggregation(TSOptionsParser &parser) { + auto &type = option_.aggregator.type; + if (parser.EatEqICase("AVG")) { + type = TSAggregatorType::AVG; + } else if (parser.EatEqICase("SUM")) { + type = TSAggregatorType::SUM; + } else if (parser.EatEqICase("MIN")) { + type = TSAggregatorType::MIN; + } else if (parser.EatEqICase("MAX")) { + type = TSAggregatorType::MAX; + } else if (parser.EatEqICase("RANGE")) { + type = TSAggregatorType::RANGE; + } else if (parser.EatEqICase("COUNT")) { + type = TSAggregatorType::COUNT; + } else if (parser.EatEqICase("FIRST")) { + type = TSAggregatorType::FIRST; + } else if (parser.EatEqICase("LAST")) { + type = TSAggregatorType::LAST; + } else if (parser.EatEqICase("STD.P")) { + type = TSAggregatorType::STD_P; + } else if (parser.EatEqICase("STD.S")) { + type = TSAggregatorType::STD_S; + } else if (parser.EatEqICase("VAR.P")) { + type = TSAggregatorType::VAR_P; + } else if (parser.EatEqICase("VAR.S")) { + type = TSAggregatorType::VAR_S; + } else { + return {Status::RedisParseErr, "Invalid aggregator type"}; + } + + auto duration = parser.TakeInt(); + if (!duration.IsOK()) { + return {Status::RedisParseErr, "Couldn't parse AGGREGATION"}; + } + option_.aggregator.bucket_duration = duration.GetValue(); + if (option_.aggregator.bucket_duration == 0) { + return {Status::RedisParseErr, "bucketDuration must be greater than zero"}; + } + return Status::OK(); + } + + Status handleBucketTimestamp(TSOptionsParser &parser) { + if (option_.aggregator.type == TSAggregatorType::NONE) { + return {Status::RedisParseErr, "BUCKETTIMESTAMP flag should be the 3rd or 4th flag after AGGREGATION flag"}; + } + using BucketTimestampType = TSRangeOption::BucketTimestampType; + if (parser.EatEqICase("START")) { + option_.bucket_timestamp_type = BucketTimestampType::Start; + } else if (parser.EatEqICase("END")) { + option_.bucket_timestamp_type = BucketTimestampType::End; + } else if (parser.EatEqICase("MID")) { + option_.bucket_timestamp_type = BucketTimestampType::Mid; + } else { + return {Status::RedisParseErr, "unknown BUCKETTIMESTAMP parameter"}; + } + return Status::OK(); + } + + Status handleEmpty([[maybe_unused]] TSOptionsParser &parser) { + if (option_.aggregator.type == TSAggregatorType::NONE) { + return {Status::RedisParseErr, "EMPTY flag should be the 3rd or 5th flag after AGGREGATION flag"}; + } + option_.is_return_empty = true; + return Status::OK(); + } +}; + +class CommandTSRange : public CommandTSRangeBase { + public: + CommandTSRange() : CommandTSRangeBase(2, 0) {} + Status Parse(const std::vector &args) override { + if (args.size() < 4) { + return {Status::RedisParseErr, "wrong number of arguments for 'ts.range' command"}; + } + + user_key_ = args[1]; + + return CommandTSRangeBase::Parse(args); + } + + Status Execute(engine::Context &ctx, Server *srv, Connection *conn, std::string *output) override { + auto timeseries_db = TimeSeries(srv->storage, conn->GetNamespace()); + std::vector res; + auto s = timeseries_db.Range(ctx, user_key_, GetRangeOption(), &res); + if (!s.ok()) return {Status::RedisExecErr, errKeyNotFound}; + std::vector reply; + reply.reserve(res.size()); + for (auto &sample : res) { + reply.push_back(FormatTSSampleAsRedisReply(sample)); + } + *output = redis::Array(reply); + return Status::OK(); + } + + private: + std::string user_key_; +}; + REDIS_REGISTER_COMMANDS(Timeseries, MakeCmdAttr("ts.create", -2, "write", 1, 1, 1), MakeCmdAttr("ts.add", -4, "write", 1, 1, 1), MakeCmdAttr("ts.madd", -4, "write", 1, -3, 1), + MakeCmdAttr("ts.range", -4, "read-only", 1, 1, 1), MakeCmdAttr("ts.info", -2, "read-only", 1, 1, 1)); } // namespace redis diff --git a/src/types/redis_timeseries.cc b/src/types/redis_timeseries.cc index eb0bdb128ca..1c4a1a702ed 100644 --- a/src/types/redis_timeseries.cc +++ b/src/types/redis_timeseries.cc @@ -31,10 +31,87 @@ constexpr uint64_t kDefaultChunkSize = 1024; constexpr auto kDefaultChunkType = TimeSeriesMetadata::ChunkType::UNCOMPRESSED; constexpr auto kDefaultDuplicatePolicy = TimeSeriesMetadata::DuplicatePolicy::BLOCK; +std::vector AggregateSamplesByRangeOption(std::vector samples, const TSRangeOption &option) { + const auto &aggregator = option.aggregator; + std::vector res; + if (aggregator.type == TSAggregatorType::NONE || samples.empty()) { + res = std::move(samples); + return res; + } + uint64_t start_bucket = aggregator.CalculateAlignedBucketLeft(samples.front().ts); + uint64_t end_bucket = aggregator.CalculateAlignedBucketLeft(samples.back().ts); + uint64_t bucket_count = (end_bucket - start_bucket) / aggregator.bucket_duration + 1; + + std::vector> spans; + spans.reserve(bucket_count); + auto it = samples.begin(); + const auto end = samples.end(); + uint64_t bucket_left = start_bucket; + while (it != end) { + uint64_t bucket_right = aggregator.CalculateAlignedBucketRight(bucket_left); + auto lower = std::lower_bound(it, end, TSSample{bucket_left, 0.0}); + auto upper = std::lower_bound(lower, end, TSSample{bucket_right, 0.0}); + spans.emplace_back(lower, upper); + it = upper; + + bucket_left = bucket_right; + } + + auto get_bucket_ts = [&](uint64_t left) -> uint64_t { + using BucketTimestampType = TSRangeOption::BucketTimestampType; + switch (option.bucket_timestamp_type) { + case BucketTimestampType::Start: + return left; + case BucketTimestampType::End: + return left + aggregator.bucket_duration; + case BucketTimestampType::Mid: + return left + aggregator.bucket_duration / 2; + default: + unreachable(); + } + return 0; + }; + res.reserve(spans.size()); + bucket_left = start_bucket; + for (size_t i = 0; i < spans.size(); i++) { + if (option.count_limit && res.size() >= option.count_limit) { + break; + } + TSSample sample; + if (i != 0) { + bucket_left = aggregator.CalculateAlignedBucketRight(bucket_left); + } + sample.ts = get_bucket_ts(bucket_left); + if (option.is_return_empty && spans[i].empty()) { + switch (aggregator.type) { + case TSAggregatorType::SUM: + case TSAggregatorType::COUNT: + sample.v = 0; + break; + case TSAggregatorType::LAST: + if (i == 0 || spans[i - 1].empty()) { + sample.v = TSSample::NAN_VALUE; + } else { + sample.v = spans[i].back().v; + } + break; + default: + sample.v = TSSample::NAN_VALUE; + } + } else if (!spans[i].empty()) { + sample.v = aggregator.AggregateSamplesValue(spans[i]); + } else { + continue; + } + res.emplace_back(sample); + } + return res; +} + void TSDownStreamMeta::Encode(std::string *dst) const { - PutFixed8(dst, static_cast(aggregator)); - PutFixed64(dst, bucket_duration); - PutFixed64(dst, alignment); + PutFixed8(dst, static_cast(aggregator.type)); + PutFixed64(dst, aggregator.bucket_duration); + PutFixed64(dst, aggregator.alignment); PutFixed64(dst, latest_bucket_idx); PutFixed8(dst, static_cast(u64_auxs.size())); PutFixed8(dst, static_cast(f64_auxs.size())); @@ -51,9 +128,9 @@ rocksdb::Status TSDownStreamMeta::Decode(Slice *input) { return rocksdb::Status::InvalidArgument("TSDownStreamMeta size is too short"); } - GetFixed8(input, reinterpret_cast(&aggregator)); - GetFixed64(input, &bucket_duration); - GetFixed64(input, &alignment); + GetFixed8(input, reinterpret_cast(&aggregator.type)); + GetFixed64(input, &aggregator.bucket_duration); + GetFixed64(input, &aggregator.alignment); GetFixed64(input, &latest_bucket_idx); uint8_t u64_auxs_size = 0; GetFixed8(input, &u64_auxs_size); @@ -114,6 +191,144 @@ TimeSeriesMetadata CreateMetadataFromOption(const TSCreateOption &option) { return metadata; } +uint64_t TSAggregator::CalculateAlignedBucketLeft(uint64_t ts) const { + uint64_t x = 0; + + if (ts >= alignment) { + uint64_t diff = ts - alignment; + uint64_t k = diff / bucket_duration; + x = alignment + k * bucket_duration; + } else { + uint64_t diff = alignment - ts; + uint64_t m0 = diff / bucket_duration + (diff % bucket_duration == 0 ? 0 : 1); + if (m0 <= alignment / bucket_duration) { + x = alignment - m0 * bucket_duration; + } + } + + return x; +} + +uint64_t TSAggregator::CalculateAlignedBucketRight(uint64_t ts) const { + uint64_t x = TSSample::MAX_TIMESTAMP; + if (ts < alignment) { + uint64_t diff = alignment - ts; + uint64_t k = diff / bucket_duration; + x = alignment - k * bucket_duration; + } else { + uint64_t diff = ts - alignment; + uint64_t m0 = diff / bucket_duration + 1; + if (m0 <= (TSSample::MAX_TIMESTAMP - alignment) / bucket_duration) { + x = alignment + m0 * bucket_duration; + } + } + + return x; +} + +double TSAggregator::AggregateSamplesValue(nonstd::span samples) const { + double res = TSSample::NAN_VALUE; + if (samples.empty()) { + return res; + } + auto sample_size = static_cast(samples.size()); + switch (type) { + case TSAggregatorType::AVG: { + res = std::accumulate(samples.begin(), samples.end(), 0.0, + [](double sum, const TSSample &sample) { return sum + sample.v; }) / + sample_size; + break; + } + case TSAggregatorType::SUM: { + res = std::accumulate(samples.begin(), samples.end(), 0.0, + [](double sum, const TSSample &sample) { return sum + sample.v; }); + break; + } + case TSAggregatorType::MIN: { + res = std::min_element(samples.begin(), samples.end(), [](const TSSample &a, const TSSample &b) { + return a.v < b.v; + })->v; + break; + } + case TSAggregatorType::MAX: { + res = std::max_element(samples.begin(), samples.end(), [](const TSSample &a, const TSSample &b) { + return a.v < b.v; + })->v; + break; + } + case TSAggregatorType::RANGE: { + auto [min_it, max_it] = std::minmax_element(samples.begin(), samples.end(), + [](const TSSample &a, const TSSample &b) { return a.v < b.v; }); + res = max_it->v - min_it->v; + break; + } + case TSAggregatorType::COUNT: { + res = sample_size; + break; + } + case TSAggregatorType::FIRST: { + res = samples.front().v; + break; + } + case TSAggregatorType::LAST: { + res = samples.back().v; + break; + } + case TSAggregatorType::STD_P: { + double mean = std::accumulate(samples.begin(), samples.end(), 0.0, + [](double sum, const TSSample &sample) { return sum + sample.v; }) / + sample_size; + double variance = + std::accumulate(samples.begin(), samples.end(), 0.0, + [mean](double sum, const TSSample &sample) { return sum + std::pow(sample.v - mean, 2); }) / + sample_size; + res = std::sqrt(variance); + break; + } + case TSAggregatorType::STD_S: { + if (samples.size() <= 1) { + res = 0.0; + break; + } + double mean = std::accumulate(samples.begin(), samples.end(), 0.0, + [](double sum, const TSSample &sample) { return sum + sample.v; }) / + sample_size; + double variance = + std::accumulate(samples.begin(), samples.end(), 0.0, + [mean](double sum, const TSSample &sample) { return sum + std::pow(sample.v - mean, 2); }) / + (sample_size - 1.0); + res = std::sqrt(variance); + break; + } + case TSAggregatorType::VAR_P: { + double mean = std::accumulate(samples.begin(), samples.end(), 0.0, + [](double sum, const TSSample &sample) { return sum + sample.v; }) / + sample_size; + res = std::accumulate(samples.begin(), samples.end(), 0.0, + [mean](double sum, const TSSample &sample) { return sum + std::pow(sample.v - mean, 2); }) / + sample_size; + break; + } + case TSAggregatorType::VAR_S: { + if (samples.size() <= 1) { + res = 0.0; + break; + } + double mean = std::accumulate(samples.begin(), samples.end(), 0.0, + [](double sum, const TSSample &sample) { return sum + sample.v; }) / + sample_size; + res = std::accumulate(samples.begin(), samples.end(), 0.0, + [mean](double sum, const TSSample &sample) { return sum + std::pow(sample.v - mean, 2); }) / + (sample_size - 1.0); + break; + } + default: + unreachable(); + } + + return res; +} + rocksdb::Status TimeSeries::getTimeSeriesMetadata(engine::Context &ctx, const Slice &ns_key, TimeSeriesMetadata *metadata) { return Database::GetMetadata(ctx, {kRedisTimeSeries}, ns_key, metadata); @@ -454,4 +669,117 @@ rocksdb::Status TimeSeries::Info(engine::Context &ctx, const Slice &user_key, TS return rocksdb::Status::OK(); } +rocksdb::Status TimeSeries::Range(engine::Context &ctx, const Slice &user_key, const TSRangeOption &option, + std::vector *res) { + std::string ns_key = AppendNamespacePrefix(user_key); + + TimeSeriesMetadata metadata(false); + rocksdb::Status s = getTimeSeriesMetadata(ctx, ns_key, &metadata); + if (!s.ok()) { + return s; + } + if (option.end_ts < option.start_ts) { + return rocksdb::Status::OK(); + } + + // In the emun `TSSubkeyType`, `LABEL` is the next of `CHUNK` + std::string chunk_upper_bound = internalKeyFromLabelKey(ns_key, metadata, ""); + std::string end_key = internalKeyFromChunkID(ns_key, metadata, TSSample::MAX_TIMESTAMP); + std::string prefix = end_key.substr(0, end_key.size() - sizeof(uint64_t)); + + rocksdb::ReadOptions read_options = ctx.DefaultScanOptions(); + rocksdb::Slice upper_bound(chunk_upper_bound); + read_options.iterate_upper_bound = &upper_bound; + rocksdb::Slice lower_bound(prefix); + read_options.iterate_lower_bound = &lower_bound; + + // Get the latest chunk + auto iter = util::UniqueIterator(ctx, read_options); + iter->SeekForPrev(end_key); + if (!iter->Valid() || !iter->key().starts_with(prefix)) { + return rocksdb::Status::OK(); + } + auto chunk = CreateTSChunkFromData(iter->value()); + uint64_t last_timestamp = chunk->GetLastTimestamp(); + uint64_t retention_bound = (metadata.retention_time == 0 || last_timestamp <= metadata.retention_time) + ? 0 + : last_timestamp - metadata.retention_time; + uint64_t start_timestamp = std::max(retention_bound, option.start_ts); + uint64_t end_timestamp = std::min(last_timestamp, option.end_ts); + + // Update iterator options + auto start_key = internalKeyFromChunkID(ns_key, metadata, start_timestamp); + if (end_timestamp != TSSample::MAX_TIMESTAMP) { + end_key = internalKeyFromChunkID(ns_key, metadata, end_timestamp + 1); + } + upper_bound = Slice(end_key); + read_options.iterate_upper_bound = &upper_bound; + iter = util::UniqueIterator(ctx, read_options); + + iter->SeekForPrev(start_key); + if (!iter->Valid()) { + iter->Seek(start_key); + } else if (!iter->key().starts_with(prefix)) { + iter->Next(); + } + // Prepare to store results + std::vector temp_results; + const auto &aggregator = option.aggregator; + bool has_aggregator = aggregator.type != TSAggregatorType::NONE; + if (iter->Valid()) { + if (option.count_limit != 0 && !has_aggregator) { + temp_results.reserve(option.count_limit); + } else { + chunk = CreateTSChunkFromData(iter->value()); + auto range = chunk->GetLastTimestamp() - chunk->GetFirstTimestamp() + 1; + auto estimate_chunks = std::min((end_timestamp - start_timestamp) / range, uint64_t(32)); + temp_results.reserve(estimate_chunks * metadata.chunk_size); + } + } + // Get samples from chunks + uint64_t bucket_count = 0; + uint64_t last_bucket = 0; + bool is_not_enough = true; + for (; iter->Valid() && is_not_enough; iter->Next()) { + chunk = CreateTSChunkFromData(iter->value()); + auto it = chunk->CreateIterator(); + while (it->HasNext()) { + auto sample = it->Next().value(); + // Early termination check + if (!has_aggregator && option.count_limit && temp_results.size() >= option.count_limit) { + is_not_enough = false; + break; + } + const bool in_time_range = sample->ts >= start_timestamp && sample->ts <= end_timestamp; + const bool not_time_filtered = option.filter_by_ts.empty() || option.filter_by_ts.count(sample->ts); + const bool value_in_range = !option.filter_by_value || (sample->v >= option.filter_by_value->first && + sample->v <= option.filter_by_value->second); + + if (!in_time_range || !not_time_filtered || !value_in_range) { + continue; + } + + // Do checks for early termination when `count_limit` is set. + if (has_aggregator && option.count_limit > 0) { + const auto bucket = aggregator.CalculateAlignedBucketRight(sample->ts); + const bool is_empty_count = (last_bucket > 0 && option.is_return_empty); + const size_t increment = is_empty_count ? (bucket - last_bucket) / aggregator.bucket_duration : 1; + bucket_count += increment; + last_bucket = bucket; + if (bucket_count > option.count_limit) { + is_not_enough = false; + temp_results.push_back(*sample); // Ensure empty bucket is reported + break; + } + } + temp_results.push_back(*sample); + } + } + + // Process compaction logic + *res = AggregateSamplesByRangeOption(std::move(temp_results), option); + + return rocksdb::Status::OK(); +} + } // namespace redis diff --git a/src/types/redis_timeseries.h b/src/types/redis_timeseries.h index 16acdacf71c..394d90ac7d9 100644 --- a/src/types/redis_timeseries.h +++ b/src/types/redis_timeseries.h @@ -40,24 +40,44 @@ enum class IndexKeyType : uint8_t { }; enum class TSAggregatorType : uint8_t { - AVG = 0, - SUM = 1, - MIN = 2, - MAX = 3, - RANGE = 4, - COUNT = 5, - FIRST = 6, - LAST = 7, - STD_P = 8, - STD_S = 9, - VAR_P = 10, - VAR_S = 11, + NONE = 0, + AVG = 1, + SUM = 2, + MIN = 3, + MAX = 4, + RANGE = 5, + COUNT = 6, + FIRST = 7, + LAST = 8, + STD_P = 9, + STD_S = 10, + VAR_P = 11, + VAR_S = 12, +}; + +struct TSAggregator { + TSAggregatorType type = TSAggregatorType::NONE; + uint64_t bucket_duration = 0; + uint64_t alignment = 0; + + TSAggregator() = default; + TSAggregator(TSAggregatorType type, uint64_t bucket_duration, uint64_t alignment) + : type(type), bucket_duration(bucket_duration), alignment(alignment) {} + + // Calculates the start timestamp of the aligned bucket that contains the given timestamp. + // E.g. `ts`=100, `duration`=30, `alignment`=20. + // The bucket containing `ts=100` starts at `80` (since 80 ≤ 100 < 110). Returns `80`. + uint64_t CalculateAlignedBucketLeft(uint64_t ts) const; + + // Calculates the end timestamp of the aligned bucket that contains the given timestamp. + uint64_t CalculateAlignedBucketRight(uint64_t ts) const; + + // Calculates the aggregated value of the given samples according to the aggregator type + double AggregateSamplesValue(nonstd::span samples) const; }; struct TSDownStreamMeta { - TSAggregatorType aggregator; - uint64_t bucket_duration; - uint64_t alignment; + TSAggregator aggregator; uint64_t latest_bucket_idx; // store auxiliary info for each aggregator. @@ -66,12 +86,8 @@ struct TSDownStreamMeta { std::vector f64_auxs; TSDownStreamMeta() = default; - TSDownStreamMeta(TSAggregatorType aggregator, uint64_t bucket_duration, uint64_t alignment, - uint64_t latest_bucket_idx) - : aggregator(aggregator), - bucket_duration(bucket_duration), - alignment(alignment), - latest_bucket_idx(latest_bucket_idx) {} + TSDownStreamMeta(TSAggregatorType agg_type, uint64_t bucket_duration, uint64_t alignment, uint64_t latest_bucket_idx) + : aggregator(agg_type, bucket_duration, alignment), latest_bucket_idx(latest_bucket_idx) {} void Encode(std::string *dst) const; rocksdb::Status Decode(Slice *input); @@ -116,6 +132,25 @@ struct TSInfoResult { LabelKVList labels; }; +struct TSRangeOption { + enum class BucketTimestampType : uint8_t { + Start = 0, + End = 1, + Mid = 2, + }; + uint64_t start_ts = 0; + uint64_t end_ts = TSSample::MAX_TIMESTAMP; + uint64_t count_limit = 0; + std::set filter_by_ts; + std::optional> filter_by_value; + + // Used for comapction + TSAggregator aggregator; + bool is_return_latest = false; + bool is_return_empty = false; + BucketTimestampType bucket_timestamp_type = BucketTimestampType::Start; +}; + TimeSeriesMetadata CreateMetadataFromOption(const TSCreateOption &option); class TimeSeries : public SubKeyScanner { @@ -131,6 +166,8 @@ class TimeSeries : public SubKeyScanner { rocksdb::Status MAdd(engine::Context &ctx, const Slice &user_key, std::vector samples, std::vector *res); rocksdb::Status Info(engine::Context &ctx, const Slice &user_key, TSInfoResult *res); + rocksdb::Status Range(engine::Context &ctx, const Slice &user_key, const TSRangeOption &option, + std::vector *res); private: rocksdb::Status getTimeSeriesMetadata(engine::Context &ctx, const Slice &ns_key, TimeSeriesMetadata *metadata); diff --git a/src/types/timeseries.h b/src/types/timeseries.h index f0e883ad28c..af418cded3c 100644 --- a/src/types/timeseries.h +++ b/src/types/timeseries.h @@ -42,6 +42,7 @@ struct TSSample { double v; static constexpr uint64_t MAX_TIMESTAMP = std::numeric_limits::max(); + static constexpr double NAN_VALUE = std::numeric_limits::quiet_NaN(); // Custom comparison operator for sorting by ts bool operator<(const TSSample& other) const { return ts < other.ts; } diff --git a/tests/cppunit/types/timeseries_test.cc b/tests/cppunit/types/timeseries_test.cc index 5366bff7eed..70a5b0423d9 100644 --- a/tests/cppunit/types/timeseries_test.cc +++ b/tests/cppunit/types/timeseries_test.cc @@ -104,3 +104,231 @@ TEST_F(TimeSeriesTest, MAdd) { EXPECT_EQ(results.size(), 1); EXPECT_EQ(results[0].first, TSChunk::AddResult::kBlock); } + +TEST_F(TimeSeriesTest, Range) { + redis::TSCreateOption option; + option.labels = {{"type", "stock"}, {"name", "A"}}; + auto s = ts_db_->Create(*ctx_, key_, option); + EXPECT_TRUE(s.ok()); + + // Add three batches of samples + std::vector samples1 = {{1000, 100}, {1010, 110}, {1020, 120}}; + std::vector results1; + results1.resize(samples1.size()); + s = ts_db_->MAdd(*ctx_, key_, samples1, &results1); + EXPECT_TRUE(s.ok()); + + std::vector samples2 = {{2000, 200}, {2010, 210}, {2020, 220}}; + std::vector results2; + results2.resize(samples2.size()); + s = ts_db_->MAdd(*ctx_, key_, samples2, &results2); + EXPECT_TRUE(s.ok()); + + std::vector samples3 = {{3000, 300}, {3010, 310}, {3020, 320}}; + std::vector results3; + results3.resize(samples3.size()); + s = ts_db_->MAdd(*ctx_, key_, samples3, &results3); + EXPECT_TRUE(s.ok()); + + // Test basic range query without aggregation + std::vector res; + redis::TSRangeOption range_opt; + range_opt.start_ts = 0; + range_opt.end_ts = TSSample::MAX_TIMESTAMP; + s = ts_db_->Range(*ctx_, key_, range_opt, &res); + EXPECT_TRUE(s.ok()); + EXPECT_EQ(res.size(), 9); + for (size_t i = 0; i < samples1.size(); ++i) { + EXPECT_EQ(res[i], samples1[i]); + } + for (size_t i = 0; i < samples2.size(); ++i) { + EXPECT_EQ(res[i + samples1.size()], samples2[i]); + } + for (size_t i = 0; i < samples3.size(); ++i) { + EXPECT_EQ(res[i + samples1.size() + samples2.size()], samples3[i]); + } + + // Test aggregation with min + res.clear(); + range_opt.aggregator.type = redis::TSAggregatorType::MIN; + range_opt.aggregator.bucket_duration = 20; + s = ts_db_->Range(*ctx_, key_, range_opt, &res); + EXPECT_TRUE(s.ok()); + EXPECT_EQ(res.size(), 6); + EXPECT_EQ(res[0].ts, 1000); + EXPECT_EQ(res[0].v, 100); + EXPECT_EQ(res[1].ts, 1020); + EXPECT_EQ(res[1].v, 120); + EXPECT_EQ(res[2].ts, 2000); + EXPECT_EQ(res[2].v, 200); + EXPECT_EQ(res[3].ts, 2020); + EXPECT_EQ(res[3].v, 220); + EXPECT_EQ(res[4].ts, 3000); + EXPECT_EQ(res[4].v, 300); + EXPECT_EQ(res[5].ts, 3020); + EXPECT_EQ(res[5].v, 320); + + // Test different aggregators + res.clear(); + range_opt.aggregator.type = redis::TSAggregatorType::AVG; + range_opt.aggregator.bucket_duration = 1000; + s = ts_db_->Range(*ctx_, key_, range_opt, &res); + EXPECT_TRUE(s.ok()); + EXPECT_EQ(res[0].v, 110); + EXPECT_EQ(res[1].v, 210); + EXPECT_EQ(res[2].v, 310); + + range_opt.aggregator.type = redis::TSAggregatorType::SUM; + s = ts_db_->Range(*ctx_, key_, range_opt, &res); + EXPECT_TRUE(s.ok()); + EXPECT_EQ(res[0].v, 330); + EXPECT_EQ(res[1].v, 630); + EXPECT_EQ(res[2].v, 930); + + range_opt.aggregator.type = redis::TSAggregatorType::COUNT; + s = ts_db_->Range(*ctx_, key_, range_opt, &res); + EXPECT_TRUE(s.ok()); + EXPECT_EQ(res[0].v, 3); + EXPECT_EQ(res[1].v, 3); + EXPECT_EQ(res[2].v, 3); + + range_opt.aggregator.type = redis::TSAggregatorType::RANGE; + s = ts_db_->Range(*ctx_, key_, range_opt, &res); + EXPECT_TRUE(s.ok()); + EXPECT_EQ(res[0].v, 20); + EXPECT_EQ(res[1].v, 20); + EXPECT_EQ(res[2].v, 20); + + range_opt.aggregator.type = redis::TSAggregatorType::FIRST; + s = ts_db_->Range(*ctx_, key_, range_opt, &res); + EXPECT_TRUE(s.ok()); + EXPECT_EQ(res[0].v, 100); + EXPECT_EQ(res[1].v, 200); + EXPECT_EQ(res[2].v, 300); + + range_opt.aggregator.type = redis::TSAggregatorType::STD_P; + s = ts_db_->Range(*ctx_, key_, range_opt, &res); + EXPECT_TRUE(s.ok()); + EXPECT_NEAR(res[0].v, 8.16496580927726, 1e-6); + EXPECT_NEAR(res[1].v, 8.16496580927726, 1e-6); + EXPECT_NEAR(res[2].v, 8.16496580927726, 1e-6); + + range_opt.aggregator.type = redis::TSAggregatorType::STD_S; + s = ts_db_->Range(*ctx_, key_, range_opt, &res); + EXPECT_TRUE(s.ok()); + EXPECT_NEAR(res[0].v, 10.0, 1e-6); + EXPECT_NEAR(res[1].v, 10.0, 1e-6); + EXPECT_NEAR(res[2].v, 10.0, 1e-6); + + range_opt.aggregator.type = redis::TSAggregatorType::VAR_P; + s = ts_db_->Range(*ctx_, key_, range_opt, &res); + EXPECT_TRUE(s.ok()); + EXPECT_NEAR(res[0].v, 66.666666, 1e-6); + EXPECT_NEAR(res[1].v, 66.666666, 1e-6); + EXPECT_NEAR(res[2].v, 66.666666, 1e-6); + + range_opt.aggregator.type = redis::TSAggregatorType::VAR_S; + s = ts_db_->Range(*ctx_, key_, range_opt, &res); + EXPECT_TRUE(s.ok()); + EXPECT_NEAR(res[0].v, 100.0, 1e-6); + EXPECT_NEAR(res[1].v, 100.0, 1e-6); + EXPECT_NEAR(res[2].v, 100.0, 1e-6); + + // Test alignment + res.clear(); + range_opt.aggregator.type = redis::TSAggregatorType::MIN; + range_opt.aggregator.alignment = 10; + range_opt.aggregator.bucket_duration = 20; + s = ts_db_->Range(*ctx_, key_, range_opt, &res); + EXPECT_TRUE(s.ok()); + EXPECT_EQ(res.size(), 6); + EXPECT_EQ(res[0].ts, 990); + EXPECT_EQ(res[0].v, 100); + EXPECT_EQ(res[1].ts, 1010); + EXPECT_EQ(res[1].v, 110); + EXPECT_EQ(res[2].ts, 1990); + EXPECT_EQ(res[2].v, 200); + EXPECT_EQ(res[3].ts, 2010); + EXPECT_EQ(res[3].v, 210); + EXPECT_EQ(res[4].ts, 2990); + EXPECT_EQ(res[4].v, 300); + EXPECT_EQ(res[5].ts, 3010); + EXPECT_EQ(res[5].v, 310); + + // Test alignment with irregular bucket + res.clear(); + range_opt.aggregator.bucket_duration = 4000; + range_opt.aggregator.alignment = 2000; + range_opt.aggregator.type = redis::TSAggregatorType::SUM; + range_opt.start_ts = 1000; + s = ts_db_->Range(*ctx_, key_, range_opt, &res); + EXPECT_TRUE(s.ok()); + EXPECT_EQ(res.size(), 2); + EXPECT_EQ(res[0].ts, 0); + EXPECT_EQ(res[0].v, 330); + EXPECT_EQ(res[1].ts, 2000); + EXPECT_EQ(res[1].v, 1560); + + // Test bucket timestamp type + res.clear(); + range_opt.aggregator.type = redis::TSAggregatorType::MIN; + range_opt.aggregator.alignment = 10; + range_opt.aggregator.bucket_duration = 20; + range_opt.start_ts = 0; + range_opt.bucket_timestamp_type = redis::TSRangeOption::BucketTimestampType::Mid; + s = ts_db_->Range(*ctx_, key_, range_opt, &res); + EXPECT_TRUE(s.ok()); + EXPECT_EQ(res.size(), 6); + EXPECT_EQ(res[0].ts, 1000); + EXPECT_EQ(res[0].v, 100); + EXPECT_EQ(res[1].ts, 1020); + EXPECT_EQ(res[1].v, 110); + EXPECT_EQ(res[2].ts, 2000); + EXPECT_EQ(res[2].v, 200); + EXPECT_EQ(res[3].ts, 2020); + EXPECT_EQ(res[3].v, 210); + EXPECT_EQ(res[4].ts, 3000); + EXPECT_EQ(res[4].v, 300); + EXPECT_EQ(res[5].ts, 3020); + EXPECT_EQ(res[5].v, 310); + + // Test empty buckets + res.clear(); + range_opt.is_return_empty = true; + range_opt.start_ts = 1500; + range_opt.end_ts = 2500; + range_opt.aggregator.bucket_duration = 5; + s = ts_db_->Range(*ctx_, key_, range_opt, &res); + EXPECT_TRUE(s.ok()); + EXPECT_EQ(res.size(), 5); + EXPECT_EQ(res[0].ts, 2002); + EXPECT_EQ(res[0].v, 200); + EXPECT_EQ(res[1].ts, 2007); + EXPECT_TRUE(std::isnan(res[1].v)); + EXPECT_EQ(res[2].ts, 2012); + EXPECT_EQ(res[2].v, 210); + EXPECT_EQ(res[3].ts, 2017); + EXPECT_TRUE(std::isnan(res[3].v)); + EXPECT_EQ(res[4].ts, 2022); + EXPECT_EQ(res[4].v, 220); + + // Test filter by value + res.clear(); + range_opt.aggregator.bucket_duration = 20; + range_opt.is_return_empty = false; + range_opt.filter_by_value = std::make_optional(std::make_pair(200.0, 300.0)); + s = ts_db_->Range(*ctx_, key_, range_opt, &res); + EXPECT_TRUE(s.ok()); + EXPECT_EQ(res.size(), 2); + for (const auto& sample : res) { + EXPECT_GE(sample.v, 200.0); + EXPECT_LE(sample.v, 300.0); + } + + // Test count limit + res.clear(); + range_opt.count_limit = 1; + s = ts_db_->Range(*ctx_, key_, range_opt, &res); + EXPECT_TRUE(s.ok()); + EXPECT_EQ(res.size(), 1); +} diff --git a/tests/gocase/unit/type/timeseries/timeseries_test.go b/tests/gocase/unit/type/timeseries/timeseries_test.go index ed23a3dfc5f..4f2fcd6c5f5 100644 --- a/tests/gocase/unit/type/timeseries/timeseries_test.go +++ b/tests/gocase/unit/type/timeseries/timeseries_test.go @@ -20,6 +20,7 @@ package timeseries import ( "context" + "math" "strconv" "testing" "time" @@ -228,4 +229,183 @@ func testTimeSeries(t *testing.T, configs util.KvrocksServerConfigs) { assert.Contains(t, res[0], "the key is not a TSDB key") assert.Equal(t, res[1], int64(1000)) }) + + t.Run("TS.RANGE Invalid Timestamp", func(t *testing.T) { + require.NoError(t, rdb.Del(ctx, key).Err()) + require.ErrorContains(t, rdb.Do(ctx, "ts.range", key, "abc", "1000").Err(), "wrong fromTimestamp") + require.ErrorContains(t, rdb.Do(ctx, "ts.range", key, "1000", "xyz").Err(), "wrong toTimestamp") + }) + + t.Run("TS.RANGE No Data", func(t *testing.T) { + require.NoError(t, rdb.Del(ctx, key).Err()) + require.NoError(t, rdb.Do(ctx, "ts.create", key).Err()) + res := rdb.Do(ctx, "ts.range", key, "-", "+").Val().([]interface{}) + assert.Empty(t, res) + }) + + t.Run("TS.RANGE Nonexistent Key", func(t *testing.T) { + require.NoError(t, rdb.Del(ctx, "nonexistent").Err()) + require.ErrorContains(t, rdb.Do(ctx, "ts.range", "nonexistent", "-", "+").Err(), "key does not exist") + }) + + t.Run("TS.RANGE Invalid Aggregation Type", func(t *testing.T) { + require.NoError(t, rdb.Del(ctx, key).Err()) + require.NoError(t, rdb.Do(ctx, "ts.create", key).Err()) + require.ErrorContains(t, rdb.Do(ctx, "ts.range", key, "-", "+", "AGGREGATION", "invalid", "1000").Err(), "Invalid aggregator type") + }) + + t.Run("TS.RANGE Invalid Aggregation Duration", func(t *testing.T) { + require.NoError(t, rdb.Del(ctx, key).Err()) + require.NoError(t, rdb.Do(ctx, "ts.create", key).Err()) + require.ErrorContains(t, rdb.Do(ctx, "ts.range", key, "-", "+", "AGGREGATION", "avg", "0").Err(), "bucketDuration must be greater than zero") + }) + + t.Run("TS.RANGE Invalid Count", func(t *testing.T) { + require.NoError(t, rdb.Del(ctx, key).Err()) + require.NoError(t, rdb.Do(ctx, "ts.create", key).Err()) + require.ErrorContains(t, rdb.Do(ctx, "ts.range", key, "-", "+", "COUNT", "0").Err(), "Invalid COUNT value") + }) + + t.Run("TS.RANGE Invalid Align Parameter", func(t *testing.T) { + require.NoError(t, rdb.Del(ctx, key).Err()) + require.NoError(t, rdb.Do(ctx, "ts.create", key).Err()) + require.ErrorContains(t, rdb.Do(ctx, "ts.range", key, "-", "+", "AGGREGATION", "avg", "1000", "ALIGN", "invalid").Err(), "unknown ALIGN parameter") + }) + + t.Run("TS.RANGE Align Without Aggregation", func(t *testing.T) { + require.NoError(t, rdb.Del(ctx, key).Err()) + require.NoError(t, rdb.Do(ctx, "ts.create", key).Err()) + require.ErrorContains(t, rdb.Do(ctx, "ts.range", key, "-", "+", "ALIGN", "1000").Err(), "ALIGN parameter can only be used with AGGREGATION") + }) + + t.Run("TS.RANGE BucketTimestamp Without Aggregation", func(t *testing.T) { + require.NoError(t, rdb.Del(ctx, key).Err()) + require.NoError(t, rdb.Do(ctx, "ts.create", key).Err()) + require.ErrorContains(t, rdb.Do(ctx, "ts.range", key, "-", "+", "BUCKETTIMESTAMP", "START").Err(), "BUCKETTIMESTAMP flag should be the 3rd or 4th flag after AGGREGATION flag") + }) + + t.Run("TS.RANGE Empty Without Aggregation", func(t *testing.T) { + require.NoError(t, rdb.Del(ctx, key).Err()) + require.NoError(t, rdb.Do(ctx, "ts.create", key).Err()) + require.ErrorContains(t, rdb.Do(ctx, "ts.range", key, "-", "+", "EMPTY").Err(), "EMPTY flag should be the 3rd or 5th flag after AGGREGATION flag") + }) + + t.Run("TS.RANGE Comprehensive Test", func(t *testing.T) { + require.NoError(t, rdb.Del(ctx, key).Err()) + require.NoError(t, rdb.Do(ctx, "ts.create", key, "labels", "type", "stock", "name", "A").Err()) + + // Add samples in three batches + samples := []struct { + ts int64 + val float64 + }{ + {1000, 100}, {1010, 110}, {1020, 120}, + {2000, 200}, {2010, 210}, {2020, 220}, + {3000, 300}, {3010, 310}, {3020, 320}, + } + for _, s := range samples { + require.Equal(t, s.ts, rdb.Do(ctx, "ts.add", key, s.ts, s.val).Val()) + } + + // Test basic range without aggregation + res := rdb.Do(ctx, "ts.range", key, "-", "+").Val().([]interface{}) + assert.Equal(t, len(samples), len(res)) + for i, s := range samples { + arr := res[i].([]interface{}) + assert.Equal(t, s.ts, arr[0]) + assert.Equal(t, s.val, arr[1]) + } + + // Test MIN aggregation with 20ms bucket + res = rdb.Do(ctx, "ts.range", key, "-", "+", "AGGREGATION", "MIN", 20).Val().([]interface{}) + assert.Equal(t, 6, len(res)) + expected := []struct { + ts int64 + val float64 + }{ + {1000, 100}, {1020, 120}, + {2000, 200}, {2020, 220}, + {3000, 300}, {3020, 320}, + } + for i, exp := range expected { + arr := res[i].([]interface{}) + assert.Equal(t, exp.ts, arr[0]) + assert.Equal(t, exp.val, arr[1]) + } + + // Test alignment with 10ms offset + res = rdb.Do(ctx, "ts.range", key, "-", "+", "AGGREGATION", "MIN", 20, "ALIGN", 10).Val().([]interface{}) + assert.Equal(t, 6, len(res)) + expected = []struct { + ts int64 + val float64 + }{ + {990, 100}, {1010, 110}, + {1990, 200}, {2010, 210}, + {2990, 300}, {3010, 310}, + } + for i, exp := range expected { + arr := res[i].([]interface{}) + assert.Equal(t, exp.ts, arr[0]) + assert.Equal(t, exp.val, arr[1]) + } + + // Test mid bucket timestamp + res = rdb.Do(ctx, "ts.range", key, "-", "+", "AGGREGATION", "MIN", 20, "ALIGN", 10, "BUCKETTIMESTAMP", "MID").Val().([]interface{}) + assert.Equal(t, 6, len(res)) + expected = []struct { + ts int64 + val float64 + }{ + {1000, 100}, {1020, 110}, + {2000, 200}, {2020, 210}, + {3000, 300}, {3020, 310}, + } + for i, exp := range expected { + arr := res[i].([]interface{}) + assert.Equal(t, exp.ts, arr[0]) + assert.Equal(t, exp.val, arr[1]) + } + + // Test empty buckets + res = rdb.Do(ctx, "ts.range", key, 1500, 2500, "AGGREGATION", "MIN", 5, "EMPTY").Val().([]interface{}) + assert.Equal(t, 5, len(res)) + expected = []struct { + ts int64 + val float64 + }{ + {2000, 200}, {2005, 0}, + {2010, 210}, {2015, 0}, + {2020, 220}, + } + for i, exp := range expected { + arr := res[i].([]interface{}) + assert.Equal(t, exp.ts, arr[0]) + if i == 1 || i == 3 { + assert.True(t, math.IsNaN(arr[1].(float64))) + } else { + assert.Equal(t, exp.val, arr[1]) + } + } + + // Test value filtering + res = rdb.Do(ctx, "ts.range", key, "-", "+", "AGGREGATION", "MIN", 20, "FILTER_BY_VALUE", 200, 300).Val().([]interface{}) + assert.Equal(t, 3, len(res)) + for _, arr := range res { + val := arr.([]interface{})[1].(float64) + assert.True(t, val >= 200 && val <= 300) + } + + // Test ts filtering + res = rdb.Do(ctx, "ts.range", key, "-", "+", "FILTER_BY_TS", "1000", "3000").Val().([]interface{}) + assert.Equal(t, 2, len(res)) + for _, arr := range res { + ts := arr.([]interface{})[0].(int64) + assert.True(t, ts == 1000 || ts == 3000) + } + + // Test count limit + res = rdb.Do(ctx, "ts.range", key, "-", "+", "AGGREGATION", "MIN", 20, "COUNT", 1).Val().([]interface{}) + assert.Equal(t, 1, len(res)) + }) } From c7ed36fb649fa6e067f17d81b2b9651d5c89d321 Mon Sep 17 00:00:00 2001 From: RX Xiao <74295100+yezhizi@users.noreply.github.com> Date: Tue, 26 Aug 2025 11:06:55 +0800 Subject: [PATCH 23/43] feat(ts): Add `TS.GET` command (#3142) Part of #3048 --- src/commands/cmd_timeseries.cc | 340 +++++++++++------- src/types/redis_timeseries.cc | 37 ++ src/types/redis_timeseries.h | 1 + src/types/timeseries.cc | 7 + src/types/timeseries.h | 4 + tests/cppunit/types/timeseries_test.cc | 32 ++ .../unit/type/timeseries/timeseries_test.go | 23 ++ 7 files changed, 314 insertions(+), 130 deletions(-) diff --git a/src/commands/cmd_timeseries.cc b/src/commands/cmd_timeseries.cc index 995c9b8728d..e9cb1cf361d 100644 --- a/src/commands/cmd_timeseries.cc +++ b/src/commands/cmd_timeseries.cc @@ -141,6 +141,8 @@ class KeywordCommandBase : public Commander { handlers_.emplace_back(keyword, std::forward(handler)); } + virtual void registerDefaultHandlers() = 0; + void setSkipNum(size_t num) { skip_num_ = num; } void setTailSkipNum(size_t num) { tail_skip_num_ = num; } @@ -154,84 +156,91 @@ class KeywordCommandBase : public Commander { class CommandTSCreateBase : public KeywordCommandBase { public: - CommandTSCreateBase(size_t skip_num, size_t tail_skip_num) : KeywordCommandBase(skip_num, tail_skip_num) { - registerHandler("RETENTION", [this](TSOptionsParser &parser) { return handleRetention(parser); }); - registerHandler("CHUNK_SIZE", [this](TSOptionsParser &parser) { return handleChunkSize(parser); }); - registerHandler("ENCODING", [this](TSOptionsParser &parser) { return handleEncoding(parser); }); - registerHandler("DUPLICATE_POLICY", [this](TSOptionsParser &parser) { return handleDuplicatePolicy(parser); }); - registerHandler("LABELS", [this](TSOptionsParser &parser) { return handleLabels(parser); }); - } + CommandTSCreateBase(size_t skip_num, size_t tail_skip_num) : KeywordCommandBase(skip_num, tail_skip_num) {} protected: const TSCreateOption &getCreateOption() const { return create_option_; } - private: - TSCreateOption create_option_; + void registerDefaultHandlers() override { + registerHandler("RETENTION", + [this](TSOptionsParser &parser) { return handleRetention(parser, create_option_.retention_time); }); + registerHandler("CHUNK_SIZE", + [this](TSOptionsParser &parser) { return handleChunkSize(parser, create_option_.chunk_size); }); + registerHandler("ENCODING", + [this](TSOptionsParser &parser) { return handleEncoding(parser, create_option_.chunk_type); }); + registerHandler("DUPLICATE_POLICY", [this](TSOptionsParser &parser) { + return handleDuplicatePolicy(parser, create_option_.duplicate_policy); + }); + registerHandler("LABELS", [this](TSOptionsParser &parser) { return handleLabels(parser, create_option_.labels); }); + } - Status handleRetention(TSOptionsParser &parser) { + static Status handleRetention(TSOptionsParser &parser, uint64_t &retention_time) { auto parse_retention = parser.TakeInt(); if (!parse_retention.IsOK()) { return {Status::RedisParseErr, errBadRetention}; } - create_option_.retention_time = parse_retention.GetValue(); + retention_time = parse_retention.GetValue(); return Status::OK(); } - Status handleChunkSize(TSOptionsParser &parser) { + static Status handleChunkSize(TSOptionsParser &parser, uint64_t &chunk_size) { auto parse_chunk_size = parser.TakeInt(); if (!parse_chunk_size.IsOK()) { return {Status::RedisParseErr, errBadChunkSize}; } - create_option_.chunk_size = parse_chunk_size.GetValue(); + chunk_size = parse_chunk_size.GetValue(); return Status::OK(); } - Status handleEncoding(TSOptionsParser &parser) { + static Status handleEncoding(TSOptionsParser &parser, ChunkType &chunk_type) { if (parser.EatEqICase("UNCOMPRESSED")) { - create_option_.chunk_type = ChunkType::UNCOMPRESSED; + chunk_type = ChunkType::UNCOMPRESSED; } else if (parser.EatEqICase("COMPRESSED")) { - create_option_.chunk_type = ChunkType::COMPRESSED; + chunk_type = ChunkType::COMPRESSED; } else { return {Status::RedisParseErr, errBadEncoding}; } return Status::OK(); } - Status handleDuplicatePolicy(TSOptionsParser &parser) { + static Status handleDuplicatePolicy(TSOptionsParser &parser, DuplicatePolicy &duplicate_policy) { if (parser.EatEqICase("BLOCK")) { - create_option_.duplicate_policy = DuplicatePolicy::BLOCK; + duplicate_policy = DuplicatePolicy::BLOCK; } else if (parser.EatEqICase("FIRST")) { - create_option_.duplicate_policy = DuplicatePolicy::FIRST; + duplicate_policy = DuplicatePolicy::FIRST; } else if (parser.EatEqICase("LAST")) { - create_option_.duplicate_policy = DuplicatePolicy::LAST; + duplicate_policy = DuplicatePolicy::LAST; } else if (parser.EatEqICase("MAX")) { - create_option_.duplicate_policy = DuplicatePolicy::MAX; + duplicate_policy = DuplicatePolicy::MAX; } else if (parser.EatEqICase("MIN")) { - create_option_.duplicate_policy = DuplicatePolicy::MIN; + duplicate_policy = DuplicatePolicy::MIN; } else if (parser.EatEqICase("SUM")) { - create_option_.duplicate_policy = DuplicatePolicy::SUM; + duplicate_policy = DuplicatePolicy::SUM; } else { return {Status::RedisParseErr, errDuplicatePolicy}; } return Status::OK(); } - Status handleLabels(TSOptionsParser &parser) { + static Status handleLabels(TSOptionsParser &parser, LabelKVList &labels) { while (parser.Good()) { auto parse_key = parser.TakeStr(); auto parse_value = parser.TakeStr(); if (!parse_key.IsOK() || !parse_value.IsOK()) { break; } - create_option_.labels.push_back({parse_key.GetValue(), parse_value.GetValue()}); + labels.push_back({parse_key.GetValue(), parse_value.GetValue()}); } return Status::OK(); } + + private: + TSCreateOption create_option_; }; class CommandTSCreate : public CommandTSCreateBase { public: - CommandTSCreate() : CommandTSCreateBase(2, 0) {} + CommandTSCreate() : CommandTSCreateBase(2, 0) { registerDefaultHandlers(); } Status Parse(const std::vector &args) override { if (args.size() < 2) { return {Status::RedisParseErr, errWrongNumOfArguments}; @@ -310,9 +319,7 @@ class CommandTSInfo : public Commander { class CommandTSAdd : public CommandTSCreateBase { public: - CommandTSAdd() : CommandTSCreateBase(4, 0) { - registerHandler("ON_DUPLICATE", [this](TSOptionsParser &parser) { return handleOnDuplicatePolicy(parser); }); - } + CommandTSAdd() : CommandTSCreateBase(4, 0) { registerDefaultHandlers(); } Status Parse(const std::vector &args) override { if (args.size() < 4) { return {Status::RedisParseErr, errWrongNumOfArguments}; @@ -345,6 +352,12 @@ class CommandTSAdd : public CommandTSCreateBase { return Status::OK(); } + protected: + void registerDefaultHandlers() override { + CommandTSCreateBase::registerDefaultHandlers(); + registerHandler("ON_DUPLICATE", [this](TSOptionsParser &parser) { return handleOnDuplicatePolicy(parser); }); + } + private: DuplicatePolicy on_dup_policy_ = DuplicatePolicy::BLOCK; bool is_on_dup_policy_set_ = false; @@ -430,20 +443,80 @@ class CommandTSMAdd : public Commander { std::unordered_map> userkey_indexes_map_; }; -class CommandTSRangeBase : public KeywordCommandBase { +class CommandTSAggregatorBase : public KeywordCommandBase { public: - CommandTSRangeBase(size_t skip_num, size_t tail_skip_num) - : KeywordCommandBase(skip_num + 2, tail_skip_num), skip_num_(skip_num) { - registerHandler("LATEST", [this](TSOptionsParser &parser) { return handleLatest(parser); }); - registerHandler("FILTER_BY_TS", [this](TSOptionsParser &parser) { return handleFilterByTS(parser); }); - registerHandler("FILTER_BY_VALUE", [this](TSOptionsParser &parser) { return handleFilterByValue(parser); }); - registerHandler("COUNT", [this](TSOptionsParser &parser) { return handleCount(parser); }); - registerHandler("ALIGN", [this](TSOptionsParser &parser) { return handleAlign(parser); }); - registerHandler("AGGREGATION", [this](TSOptionsParser &parser) { return handleAggregation(parser); }); - registerHandler("BUCKETTIMESTAMP", [this](TSOptionsParser &parser) { return handleBucketTimestamp(parser); }); - registerHandler("EMPTY", [this](TSOptionsParser &parser) { return handleEmpty(parser); }); + CommandTSAggregatorBase(size_t skip_num, size_t tail_skip_num) : KeywordCommandBase(skip_num, tail_skip_num) {} + + protected: + const TSAggregator &getAggregator() const { return aggregator_; } + + void registerDefaultHandlers() override { + registerHandler("AGGREGATION", [this](TSOptionsParser &parser) { return handleAggregation(parser, aggregator_); }); + registerHandler("ALIGN", [this](TSOptionsParser &parser) { return handleAlignCommon(parser, aggregator_); }); } + static Status handleAggregation(TSOptionsParser &parser, TSAggregator &aggregator) { + auto &type = aggregator.type; + if (parser.EatEqICase("AVG")) { + type = TSAggregatorType::AVG; + } else if (parser.EatEqICase("SUM")) { + type = TSAggregatorType::SUM; + } else if (parser.EatEqICase("MIN")) { + type = TSAggregatorType::MIN; + } else if (parser.EatEqICase("MAX")) { + type = TSAggregatorType::MAX; + } else if (parser.EatEqICase("RANGE")) { + type = TSAggregatorType::RANGE; + } else if (parser.EatEqICase("COUNT")) { + type = TSAggregatorType::COUNT; + } else if (parser.EatEqICase("FIRST")) { + type = TSAggregatorType::FIRST; + } else if (parser.EatEqICase("LAST")) { + type = TSAggregatorType::LAST; + } else if (parser.EatEqICase("STD.P")) { + type = TSAggregatorType::STD_P; + } else if (parser.EatEqICase("STD.S")) { + type = TSAggregatorType::STD_S; + } else if (parser.EatEqICase("VAR.P")) { + type = TSAggregatorType::VAR_P; + } else if (parser.EatEqICase("VAR.S")) { + type = TSAggregatorType::VAR_S; + } else { + return {Status::RedisParseErr, "Invalid aggregator type"}; + } + + auto duration = parser.TakeInt(); + if (!duration.IsOK()) { + return {Status::RedisParseErr, "Couldn't parse AGGREGATION"}; + } + aggregator.bucket_duration = duration.GetValue(); + if (aggregator.bucket_duration == 0) { + return {Status::RedisParseErr, "bucketDuration must be greater than zero"}; + } + return Status::OK(); + } + static Status handleAlignCommon(TSOptionsParser &parser, TSAggregator &aggregator) { + auto align = parser.TakeInt(); + if (!align.IsOK()) { + return {Status::RedisParseErr, errTSInvalidAlign}; + } + aggregator.alignment = align.GetValue(); + return Status::OK(); + } + static Status handleLatest([[maybe_unused]] TSOptionsParser &parser, bool &is_return_latest) { + is_return_latest = true; + return Status::OK(); + } + + private: + TSAggregator aggregator_; +}; + +class CommandTSRangeBase : public CommandTSAggregatorBase { + public: + CommandTSRangeBase(size_t skip_num, size_t tail_skip_num) + : CommandTSAggregatorBase(skip_num + 2, tail_skip_num), skip_num_(skip_num) {} + Status Parse(const std::vector &args) override { TSOptionsParser parser(std::next(args.begin(), static_cast(skip_num_)), args.end()); // Parse start timestamp @@ -480,57 +553,94 @@ class CommandTSRangeBase : public KeywordCommandBase { return s; } - const TSRangeOption &GetRangeOption() const { return option_; } - - private: - TSRangeOption option_; - size_t skip_num_; - bool is_start_explicit_set_ = false; - bool is_end_explicit_set_ = false; - bool is_alignment_explicit_set_ = false; - - Status handleLatest([[maybe_unused]] TSOptionsParser &parser) { - option_.is_return_latest = true; - return Status::OK(); + protected: + const TSRangeOption &getRangeOption() const { return option_; } + + void registerDefaultHandlers() override { + registerHandler("LATEST", + [this](TSOptionsParser &parser) { return handleLatest(parser, option_.is_return_latest); }); + registerHandler("FILTER_BY_TS", + [this](TSOptionsParser &parser) { return handleFilterByTS(parser, option_.filter_by_ts); }); + registerHandler("FILTER_BY_VALUE", + [this](TSOptionsParser &parser) { return handleFilterByValue(parser, option_.filter_by_value); }); + registerHandler("COUNT", [this](TSOptionsParser &parser) { return handleCount(parser, option_.count_limit); }); + registerHandler("ALIGN", [this](TSOptionsParser &parser) { return handleAlign(parser); }); + registerHandler("AGGREGATION", + [this](TSOptionsParser &parser) { return handleAggregation(parser, option_.aggregator); }); + registerHandler("BUCKETTIMESTAMP", + [this](TSOptionsParser &parser) { return handleBucketTimestamp(parser, option_); }); + registerHandler("EMPTY", [this](TSOptionsParser &parser) { return handleEmpty(parser, option_); }); } - Status handleFilterByTS(TSOptionsParser &parser) { - option_.filter_by_ts.clear(); + static Status handleFilterByTS(TSOptionsParser &parser, std::set &filter_by_ts) { + filter_by_ts.clear(); while (parser.Good()) { auto ts = parser.TakeInt(); if (!ts.IsOK()) break; - option_.filter_by_ts.insert(ts.GetValue()); + filter_by_ts.insert(ts.GetValue()); } return Status::OK(); } - Status handleFilterByValue(TSOptionsParser &parser) { + static Status handleFilterByValue(TSOptionsParser &parser, + std::optional> &filter_by_value) { auto min = parser.TakeFloat(); auto max = parser.TakeFloat(); if (!min.IsOK() || !max.IsOK()) { return {Status::RedisParseErr, "Invalid min or max value"}; } - option_.filter_by_value = std::make_optional(std::make_pair(min.GetValue(), max.GetValue())); + filter_by_value = std::make_optional(std::make_pair(min.GetValue(), max.GetValue())); return Status::OK(); } - Status handleCount(TSOptionsParser &parser) { + static Status handleCount(TSOptionsParser &parser, uint64_t &count_limit) { auto count = parser.TakeInt(); if (!count.IsOK()) { return {Status::RedisParseErr, "Couldn't parse COUNT"}; } - option_.count_limit = count.GetValue(); - if (option_.count_limit == 0) { + count_limit = count.GetValue(); + if (count_limit == 0) { return {Status::RedisParseErr, "Invalid COUNT value"}; } return Status::OK(); } + static Status handleBucketTimestamp(TSOptionsParser &parser, TSRangeOption &option) { + if (option.aggregator.type == TSAggregatorType::NONE) { + return {Status::RedisParseErr, "BUCKETTIMESTAMP flag should be the 3rd or 4th flag after AGGREGATION flag"}; + } + using BucketTimestampType = TSRangeOption::BucketTimestampType; + if (parser.EatEqICase("START")) { + option.bucket_timestamp_type = BucketTimestampType::Start; + } else if (parser.EatEqICase("END")) { + option.bucket_timestamp_type = BucketTimestampType::End; + } else if (parser.EatEqICase("MID")) { + option.bucket_timestamp_type = BucketTimestampType::Mid; + } else { + return {Status::RedisParseErr, "unknown BUCKETTIMESTAMP parameter"}; + } + return Status::OK(); + } + + static Status handleEmpty([[maybe_unused]] TSOptionsParser &parser, TSRangeOption &option) { + if (option.aggregator.type == TSAggregatorType::NONE) { + return {Status::RedisParseErr, "EMPTY flag should be the 3rd or 5th flag after AGGREGATION flag"}; + } + option.is_return_empty = true; + return Status::OK(); + } + + private: + TSRangeOption option_; + size_t skip_num_; + bool is_start_explicit_set_ = false; + bool is_end_explicit_set_ = false; + bool is_alignment_explicit_set_ = false; + Status handleAlign(TSOptionsParser &parser) { - auto align = parser.TakeInt(); - if (align.IsOK()) { + auto s = handleAlignCommon(parser, option_.aggregator); + if (s.IsOK()) { is_alignment_explicit_set_ = true; - option_.aggregator.alignment = align.GetValue(); return Status::OK(); } @@ -556,92 +666,55 @@ class CommandTSRangeBase : public KeywordCommandBase { is_alignment_explicit_set_ = true; return Status::OK(); } +}; - Status handleAggregation(TSOptionsParser &parser) { - auto &type = option_.aggregator.type; - if (parser.EatEqICase("AVG")) { - type = TSAggregatorType::AVG; - } else if (parser.EatEqICase("SUM")) { - type = TSAggregatorType::SUM; - } else if (parser.EatEqICase("MIN")) { - type = TSAggregatorType::MIN; - } else if (parser.EatEqICase("MAX")) { - type = TSAggregatorType::MAX; - } else if (parser.EatEqICase("RANGE")) { - type = TSAggregatorType::RANGE; - } else if (parser.EatEqICase("COUNT")) { - type = TSAggregatorType::COUNT; - } else if (parser.EatEqICase("FIRST")) { - type = TSAggregatorType::FIRST; - } else if (parser.EatEqICase("LAST")) { - type = TSAggregatorType::LAST; - } else if (parser.EatEqICase("STD.P")) { - type = TSAggregatorType::STD_P; - } else if (parser.EatEqICase("STD.S")) { - type = TSAggregatorType::STD_S; - } else if (parser.EatEqICase("VAR.P")) { - type = TSAggregatorType::VAR_P; - } else if (parser.EatEqICase("VAR.S")) { - type = TSAggregatorType::VAR_S; - } else { - return {Status::RedisParseErr, "Invalid aggregator type"}; +class CommandTSRange : public CommandTSRangeBase { + public: + CommandTSRange() : CommandTSRangeBase(2, 0) { registerDefaultHandlers(); } + Status Parse(const std::vector &args) override { + if (args.size() < 4) { + return {Status::RedisParseErr, "wrong number of arguments for 'ts.range' command"}; } - auto duration = parser.TakeInt(); - if (!duration.IsOK()) { - return {Status::RedisParseErr, "Couldn't parse AGGREGATION"}; - } - option_.aggregator.bucket_duration = duration.GetValue(); - if (option_.aggregator.bucket_duration == 0) { - return {Status::RedisParseErr, "bucketDuration must be greater than zero"}; - } - return Status::OK(); - } + user_key_ = args[1]; - Status handleBucketTimestamp(TSOptionsParser &parser) { - if (option_.aggregator.type == TSAggregatorType::NONE) { - return {Status::RedisParseErr, "BUCKETTIMESTAMP flag should be the 3rd or 4th flag after AGGREGATION flag"}; - } - using BucketTimestampType = TSRangeOption::BucketTimestampType; - if (parser.EatEqICase("START")) { - option_.bucket_timestamp_type = BucketTimestampType::Start; - } else if (parser.EatEqICase("END")) { - option_.bucket_timestamp_type = BucketTimestampType::End; - } else if (parser.EatEqICase("MID")) { - option_.bucket_timestamp_type = BucketTimestampType::Mid; - } else { - return {Status::RedisParseErr, "unknown BUCKETTIMESTAMP parameter"}; - } - return Status::OK(); + return CommandTSRangeBase::Parse(args); } - Status handleEmpty([[maybe_unused]] TSOptionsParser &parser) { - if (option_.aggregator.type == TSAggregatorType::NONE) { - return {Status::RedisParseErr, "EMPTY flag should be the 3rd or 5th flag after AGGREGATION flag"}; + Status Execute(engine::Context &ctx, Server *srv, Connection *conn, std::string *output) override { + auto timeseries_db = TimeSeries(srv->storage, conn->GetNamespace()); + std::vector res; + auto s = timeseries_db.Range(ctx, user_key_, getRangeOption(), &res); + if (!s.ok()) return {Status::RedisExecErr, errKeyNotFound}; + std::vector reply; + reply.reserve(res.size()); + for (auto &sample : res) { + reply.push_back(FormatTSSampleAsRedisReply(sample)); } - option_.is_return_empty = true; + *output = redis::Array(reply); return Status::OK(); } + + private: + std::string user_key_; }; -class CommandTSRange : public CommandTSRangeBase { +class CommandTSGet : public CommandTSAggregatorBase { public: - CommandTSRange() : CommandTSRangeBase(2, 0) {} + CommandTSGet() : CommandTSAggregatorBase(2, 0) { registerDefaultHandlers(); } Status Parse(const std::vector &args) override { - if (args.size() < 4) { - return {Status::RedisParseErr, "wrong number of arguments for 'ts.range' command"}; + if (args.size() < 2) { + return {Status::RedisParseErr, "wrong number of arguments for 'ts.get' command"}; } - user_key_ = args[1]; - - return CommandTSRangeBase::Parse(args); + return CommandTSAggregatorBase::Parse(args); } - Status Execute(engine::Context &ctx, Server *srv, Connection *conn, std::string *output) override { auto timeseries_db = TimeSeries(srv->storage, conn->GetNamespace()); std::vector res; - auto s = timeseries_db.Range(ctx, user_key_, GetRangeOption(), &res); + auto s = timeseries_db.Get(ctx, user_key_, is_return_latest_, &res); if (!s.ok()) return {Status::RedisExecErr, errKeyNotFound}; + std::vector reply; reply.reserve(res.size()); for (auto &sample : res) { @@ -651,7 +724,13 @@ class CommandTSRange : public CommandTSRangeBase { return Status::OK(); } + protected: + void registerDefaultHandlers() override { + registerHandler("LATEST", [this](TSOptionsParser &parser) { return handleLatest(parser, is_return_latest_); }); + } + private: + bool is_return_latest_ = false; std::string user_key_; }; @@ -659,6 +738,7 @@ REDIS_REGISTER_COMMANDS(Timeseries, MakeCmdAttr("ts.create", -2 MakeCmdAttr("ts.add", -4, "write", 1, 1, 1), MakeCmdAttr("ts.madd", -4, "write", 1, -3, 1), MakeCmdAttr("ts.range", -4, "read-only", 1, 1, 1), - MakeCmdAttr("ts.info", -2, "read-only", 1, 1, 1)); + MakeCmdAttr("ts.info", -2, "read-only", 1, 1, 1), + MakeCmdAttr("ts.get", -2, "read-only", 1, 1, 1)); } // namespace redis diff --git a/src/types/redis_timeseries.cc b/src/types/redis_timeseries.cc index 1c4a1a702ed..5df924f2800 100644 --- a/src/types/redis_timeseries.cc +++ b/src/types/redis_timeseries.cc @@ -782,4 +782,41 @@ rocksdb::Status TimeSeries::Range(engine::Context &ctx, const Slice &user_key, c return rocksdb::Status::OK(); } +rocksdb::Status TimeSeries::Get(engine::Context &ctx, const Slice &user_key, bool is_return_latest, + std::vector *res) { + res->clear(); + std::string ns_key = AppendNamespacePrefix(user_key); + + TimeSeriesMetadata metadata(false); + rocksdb::Status s = getTimeSeriesMetadata(ctx, ns_key, &metadata); + if (!s.ok()) { + return s; + } + + // In the emun `TSSubkeyType`, `LABEL` is the next of `CHUNK` + std::string chunk_upper_bound = internalKeyFromLabelKey(ns_key, metadata, ""); + std::string end_key = internalKeyFromChunkID(ns_key, metadata, TSSample::MAX_TIMESTAMP); + std::string prefix = end_key.substr(0, end_key.size() - sizeof(uint64_t)); + + rocksdb::ReadOptions read_options = ctx.DefaultScanOptions(); + rocksdb::Slice upper_bound(chunk_upper_bound); + read_options.iterate_upper_bound = &upper_bound; + rocksdb::Slice lower_bound(prefix); + read_options.iterate_lower_bound = &lower_bound; + + // Get the latest chunk + auto iter = util::UniqueIterator(ctx, read_options); + iter->SeekForPrev(end_key); + if (!iter->Valid() || !iter->key().starts_with(prefix)) { + return rocksdb::Status::OK(); + } + auto chunk = CreateTSChunkFromData(iter->value()); + + if (is_return_latest) { + // TODO: need process `latest` option + } + res->push_back(chunk->GetLatestSample(0)); + return rocksdb::Status::OK(); +} + } // namespace redis diff --git a/src/types/redis_timeseries.h b/src/types/redis_timeseries.h index 394d90ac7d9..4663ea375cf 100644 --- a/src/types/redis_timeseries.h +++ b/src/types/redis_timeseries.h @@ -168,6 +168,7 @@ class TimeSeries : public SubKeyScanner { rocksdb::Status Info(engine::Context &ctx, const Slice &user_key, TSInfoResult *res); rocksdb::Status Range(engine::Context &ctx, const Slice &user_key, const TSRangeOption &option, std::vector *res); + rocksdb::Status Get(engine::Context &ctx, const Slice &user_key, bool is_return_latest, std::vector *res); private: rocksdb::Status getTimeSeriesMetadata(engine::Context &ctx, const Slice &ns_key, TimeSeriesMetadata *metadata); diff --git a/src/types/timeseries.cc b/src/types/timeseries.cc index 5fb505dcc73..bb0a857e31a 100644 --- a/src/types/timeseries.cc +++ b/src/types/timeseries.cc @@ -527,3 +527,10 @@ std::string UncompTSChunk::UpdateSampleValue(uint64_t ts, double value, bool is_ return new_buffer; } + +TSSample UncompTSChunk::GetLatestSample(uint32_t idx) const { + if (metadata_.count == 0 || idx >= metadata_.count) { + unreachable(); + } + return samples_[metadata_.count - 1 - idx]; +} diff --git a/src/types/timeseries.h b/src/types/timeseries.h index af418cded3c..7302f73dc36 100644 --- a/src/types/timeseries.h +++ b/src/types/timeseries.h @@ -184,6 +184,9 @@ class TSChunk { // Returns empty string if no changes virtual std::string UpdateSampleValue(uint64_t ts, double value, bool is_add_on) const = 0; + // Get idx-th latest sample, idx=0 means latest sample + virtual TSSample GetLatestSample(uint32_t idx) const = 0; + protected: nonstd::span data_; MetaData metadata_; @@ -202,6 +205,7 @@ class UncompTSChunk : public TSChunk { bool is_fix_split_mode) const override; std::string RemoveSamplesBetween(uint64_t from, uint64_t to) const override; std::string UpdateSampleValue(uint64_t ts, double value, bool is_add_on) const override; + TSSample GetLatestSample(uint32_t idx) const override; private: nonstd::span samples_; diff --git a/tests/cppunit/types/timeseries_test.cc b/tests/cppunit/types/timeseries_test.cc index 70a5b0423d9..7939db7580b 100644 --- a/tests/cppunit/types/timeseries_test.cc +++ b/tests/cppunit/types/timeseries_test.cc @@ -332,3 +332,35 @@ TEST_F(TimeSeriesTest, Range) { EXPECT_TRUE(s.ok()); EXPECT_EQ(res.size(), 1); } + +TEST_F(TimeSeriesTest, Get) { + redis::TSCreateOption option; + auto s = ts_db_->Create(*ctx_, key_, option); + EXPECT_TRUE(s.ok()); + + std::vector res; + // Test empty timeseries + s = ts_db_->Get(*ctx_, key_, false, &res); + EXPECT_TRUE(s.ok()); + EXPECT_EQ(res.size(), 0); + + // Add multiple samples + std::vector samples = {{1, 10}, {2, 20}, {3, 30}}; + std::vector results; + results.resize(samples.size()); + + s = ts_db_->MAdd(*ctx_, key_, samples, &results); + EXPECT_TRUE(s.ok()); + + // Test basic GET (returns latest sample) + s = ts_db_->Get(*ctx_, key_, false, &res); + EXPECT_TRUE(s.ok()); + EXPECT_EQ(res.size(), 1); + EXPECT_EQ(res[0].ts, 3); + EXPECT_EQ(res[0].v, 30); + + // Test GET with empty timeseries + std::vector empty_res; + s = ts_db_->Get(*ctx_, "nonexistent_key", false, &empty_res); + EXPECT_FALSE(s.ok()); +} diff --git a/tests/gocase/unit/type/timeseries/timeseries_test.go b/tests/gocase/unit/type/timeseries/timeseries_test.go index 4f2fcd6c5f5..c7fdcf5764d 100644 --- a/tests/gocase/unit/type/timeseries/timeseries_test.go +++ b/tests/gocase/unit/type/timeseries/timeseries_test.go @@ -408,4 +408,27 @@ func testTimeSeries(t *testing.T, configs util.KvrocksServerConfigs) { res = rdb.Do(ctx, "ts.range", key, "-", "+", "AGGREGATION", "MIN", 20, "COUNT", 1).Val().([]interface{}) assert.Equal(t, 1, len(res)) }) + + t.Run("TS.GET Basic", func(t *testing.T) { + key := "test_get_key" + require.NoError(t, rdb.Del(ctx, key).Err()) + require.NoError(t, rdb.Do(ctx, "ts.create", key).Err()) + // Test GET on empty timeseries + res := rdb.Do(ctx, "ts.get", key).Val().([]interface{}) + require.Equal(t, 0, len(res)) + + // Add samples + require.Equal(t, int64(1000), rdb.Do(ctx, "ts.add", key, "1000", "12.3").Val()) + require.Equal(t, int64(2000), rdb.Do(ctx, "ts.add", key, "2000", "15.6").Val()) + + // Test basic GET + res = rdb.Do(ctx, "ts.get", key).Val().([]interface{}) + require.Equal(t, 1, len(res)) + require.Equal(t, int64(2000), res[0].([]interface{})[0]) + require.Equal(t, 15.6, res[0].([]interface{})[1]) + + // Test GET on non-existent key + _, err := rdb.Do(ctx, "ts.get", "nonexistent_key").Result() + require.ErrorContains(t, err, "key does not exist") + }) } From 3a898fe22894894e50d13e9d41ccc83c8ff3aa75 Mon Sep 17 00:00:00 2001 From: Jonah Gao Date: Tue, 26 Aug 2025 19:31:22 +0800 Subject: [PATCH 24/43] chore(config): enable `level_compaction_dynamic_level_bytes` by default (#3144) To align with the description in kvrocks.conf. https://github.com/apache/kvrocks/blob/9446fbde325767f7b64896c4539035402d0f8f7a/kvrocks.conf#L1016-L1017 --- src/config/config.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/config/config.cc b/src/config/config.cc index fb0f68f6770..746932dbec3 100644 --- a/src/config/config.cc +++ b/src/config/config.cc @@ -299,7 +299,7 @@ Config::Config() { {"rocksdb.max_bytes_for_level_multiplier", false, new IntField(&rocks_db.max_bytes_for_level_multiplier, 10, 1, 100)}, {"rocksdb.level_compaction_dynamic_level_bytes", false, - new YesNoField(&rocks_db.level_compaction_dynamic_level_bytes, false)}, + new YesNoField(&rocks_db.level_compaction_dynamic_level_bytes, true)}, {"rocksdb.max_background_jobs", false, new IntField(&rocks_db.max_background_jobs, 4, 0, 32)}, {"rocksdb.rate_limiter_auto_tuned", true, new YesNoField(&rocks_db.rate_limiter_auto_tuned, true)}, {"rocksdb.avoid_unnecessary_blocking_io", true, new YesNoField(&rocks_db.avoid_unnecessary_blocking_io, true)}, From 371157814aeec7850848bdb6c8d676031e2bf497 Mon Sep 17 00:00:00 2001 From: Jonah Gao Date: Wed, 27 Aug 2025 18:11:28 +0800 Subject: [PATCH 25/43] perf(storage): eliminate unnecessary `rocksdb::DB::ListColumnFamilies()` (#3145) Historically, it was used for doing some checks. https://github.com/apache/kvrocks/blob/2bbfe5aa9531f6e76bfd10a2cc450b9bfa0f15d9/src/storage.cc#L175-L182 But these checks no longer exist today, so this operation should be unnecessary. `ListColumnFamilies()` needs to iterate through the MANIFEST, and might be costly for a large one. For example, 80MB MANIFEST, took 1.6 seconds. --- src/storage/storage.cc | 4 ---- 1 file changed, 4 deletions(-) diff --git a/src/storage/storage.cc b/src/storage/storage.cc index 551b0ec3b1d..2f204c5060d 100644 --- a/src/storage/storage.cc +++ b/src/storage/storage.cc @@ -365,10 +365,6 @@ Status Storage::Open(DBOpenMode mode) { column_families.emplace_back(std::string(kStreamColumnFamilyName), subkey_opts); column_families.emplace_back(std::string(kSearchColumnFamilyName), search_opts); - std::vector old_column_families; - auto s = rocksdb::DB::ListColumnFamilies(options, config_->db_dir, &old_column_families); - if (!s.ok()) return {Status::NotOK, s.ToString()}; - auto start = std::chrono::high_resolution_clock::now(); switch (mode) { case DBOpenMode::kDBOpenModeDefault: { From 4b4f6845eba8475622121e94c5cdb16c0645db3d Mon Sep 17 00:00:00 2001 From: sryan yuan Date: Wed, 27 Aug 2025 20:33:32 +0800 Subject: [PATCH 26/43] fix(scan): pattern-based SCAN iterations may skip remaining keys (#3146) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pattern-based SCAN iterations may skip remaining keys in a hash slot, causing incomplete key retrieval. **Bug Reproduction:** - 10 keys in the same hash slot: ["119483", "166988", "210695", "223656", "48063", "59022", "65976", "74937", "88379", "99338"] - Initial SCAN with pattern `2*`: - Returns cursor C1 and **empty keyset** (no keys match `2*`) - Records "119483" as last scanned key - Subsequent SCAN with cursor C1 and same pattern: 1. RocksDB iterator seeks to "119483" 2. Calls `Next()` → gets "166988" (next key in slot) 3. "166988" ∉ `2*` pattern → no key returned 4. **Error**: Scan incorrectly increments slot index 5. **Result**: Remaining 8 keys in slot are skipped **Bug Fix Implementation:** When scanning with a previous scan cursor and match pattern: 1. If the last scanned key is lexicographically before the pattern's start range: → Use the pattern's minimum matching key as the seek key → Instead of using the last scanned key **Example:** - Pattern: `2*` → Minimum matching key = `"2"` (hex: \x32) - Last scanned key: `"119483"` (hex: \x31\x31...) - Since `"119483"` < `"2"` lexically: ✓ **Correct:** Seek to `"2"` ✗ **Buggy:** Seek to `"119483"` --------- Co-authored-by: Twice --- src/storage/redis_db.cc | 11 ++++++++--- tests/gocase/unit/scan/scan_test.go | 29 +++++++++++++++++++++++++++++ 2 files changed, 37 insertions(+), 3 deletions(-) diff --git a/src/storage/redis_db.cc b/src/storage/redis_db.cc index 1ff52c8b008..6f12f5890cc 100644 --- a/src/storage/redis_db.cc +++ b/src/storage/redis_db.cc @@ -339,9 +339,14 @@ rocksdb::Status Database::Scan(engine::Context &ctx, const std::string &cursor, } if (!cursor.empty()) { - iter->Seek(ns_cursor); - if (iter->Valid()) { - iter->Next(); + if (storage_->IsSlotIdEncoded() && !ns_prefix.empty() && + metadata_cf_handle_->GetComparator()->Compare(rocksdb::Slice(ns_prefix), rocksdb::Slice(ns_cursor)) > 0) { + iter->Seek(ns_prefix); + } else { + iter->Seek(ns_cursor); + if (iter->Valid()) { + iter->Next(); + } } } else if (ns_prefix.empty()) { iter->SeekToFirst(); diff --git a/tests/gocase/unit/scan/scan_test.go b/tests/gocase/unit/scan/scan_test.go index 5d2f9fb9a1f..5f7f62bea89 100644 --- a/tests/gocase/unit/scan/scan_test.go +++ b/tests/gocase/unit/scan/scan_test.go @@ -32,6 +32,35 @@ import ( "golang.org/x/exp/slices" ) +func TestScanSlotRemainingKeys(t *testing.T) { + srv := util.StartServer(t, map[string]string{"cluster-enabled": "yes"}) + defer srv.Close() + + ctx := context.Background() + rdb := srv.NewClient() + defer func() { require.NoError(t, rdb.Close()) }() + + nodeID := "YYYYYYYYYYYYYYYYYYYYYYYYYYYYYYYYYYYYYYYY" + require.NoError(t, rdb.Do(ctx, "clusterx", "SETNODEID", nodeID).Err()) + clusterNodes := fmt.Sprintf("%s %s %d master - 0-16383", nodeID, srv.Host(), srv.Port()) + require.NoError(t, rdb.Do(ctx, "clusterx", "SETNODES", clusterNodes, "1").Err()) + + slot100Keys := []string{"119483", "166988", "210695", "223656", "48063", "59022", "65976", "74937", "88379", "99338"} + for _, key := range slot100Keys { + require.NoError(t, rdb.Set(ctx, key, "1", 0).Err()) + } + require.Equal(t, slot100Keys, scanAll(t, rdb)) + + cursor, keys := scan(t, rdb, "0", "match", "2*", "count", 1) + require.Equal(t, []string(nil), keys) + cursor, keys = scan(t, rdb, cursor, "match", "2*", "count", 1) + require.Equal(t, []string{"210695"}, keys) + cursor, keys = scan(t, rdb, cursor, "match", "2*", "count", 1) + require.Equal(t, []string{"223656"}, keys) + cursor, _ = scan(t, rdb, cursor, "match", "2*", "count", 1) + require.Equal(t, "0", cursor) +} + func TestScanEmptyKey(t *testing.T) { srv := util.StartServer(t, map[string]string{}) defer srv.Close() From bd268b4d3612ceb34b9f8b9df285f136d6180acc Mon Sep 17 00:00:00 2001 From: donghao526 Date: Thu, 28 Aug 2025 12:57:09 +0800 Subject: [PATCH 27/43] style: add some comments on TDigestRank --- src/types/tdigest.h | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/src/types/tdigest.h b/src/types/tdigest.h index b096187ba29..329ed519391 100644 --- a/src/types/tdigest.h +++ b/src/types/tdigest.h @@ -162,7 +162,6 @@ inline Status TDigestRank(TD&& td, const std::vector& inputs, std::vecto std::sort(indices.begin(), indices.end(), [&inputs](size_t a, size_t b) { return inputs[a] < inputs[b]; }); result.resize(inputs.size()); - size_t i = indices.size(); double cumulative_weight = 0; @@ -176,13 +175,17 @@ inline Status TDigestRank(TD&& td, const std::vector& inputs, std::vecto auto iter = td.End(); while (i > 0) { auto centroid = GET_OR_RET(iter->GetCentroid()); + if (centroid.mean > inputs[indices[i - 1]]) { + // mean > input, accumulate weight and move to prev centroid cumulative_weight += centroid.weight; } else if (centroid.mean == inputs[indices[i - 1]]) { - auto current_mean_cumulative_weight = cumulative_weight + centroid.weight / 2; + // mean == input, calculate reverse rank with half weight of current centroid cumulative_weight += centroid.weight; auto current_mean = centroid.mean; - // cumulative all the centroids which has the same mean + auto current_mean_cumulative_weight = cumulative_weight + centroid.weight / 2; + + // handle all the prev centroids which has the same mean while (!iter->IsAtBegin() && iter->Prev()) { auto next_centroid = GET_OR_RET(iter->GetCentroid()); if (current_mean != next_centroid.mean) { @@ -194,16 +197,21 @@ inline Status TDigestRank(TD&& td, const std::vector& inputs, std::vecto cumulative_weight += centroid.weight; } + // assign the reverse rank for the inputs[indices[i - 1]] result[indices[i - 1]] = static_cast(current_mean_cumulative_weight); i--; + + // handle the prev inputs which has the same value while ((i > 0) && (inputs[indices[i]] == inputs[indices[i - 1]])) { result[indices[i - 1]] = result[indices[i]]; i--; } } else { + // mean < input, calculate reverse rank result[indices[i - 1]] = static_cast(cumulative_weight); i--; } + if (iter->IsAtBegin()) { break; } From 666224035380cea5e02578ebaa77fea10388caa5 Mon Sep 17 00:00:00 2001 From: donghao526 Date: Thu, 28 Aug 2025 13:13:43 +0800 Subject: [PATCH 28/43] refactor: remove commented code --- src/types/tdigest.h | 3 --- 1 file changed, 3 deletions(-) diff --git a/src/types/tdigest.h b/src/types/tdigest.h index 329ed519391..e7c637ce54f 100644 --- a/src/types/tdigest.h +++ b/src/types/tdigest.h @@ -81,9 +81,6 @@ class TDSample { // reference: // https://github.com/apache/arrow/blob/27bbd593625122a4a25d9471c8aaf5df54a6dcf9/cpp/src/arrow/util/tdigest.cc#L38 static inline double Lerp(double a, double b, double t) { return a + t * (b - a); } -// static inline int CalculateRank(int total_weight, double cumulative_weight, bool reverse) { -// return reverse ? total_weight - 1 - static_cast(cumulative_weight) : static_cast(cumulative_weight); -// } template inline StatusOr TDigestQuantile(TD&& td, double q) { From e7f06a276ca06e1e89f4ee06181c4d2f93e03075 Mon Sep 17 00:00:00 2001 From: donghao526 Date: Thu, 28 Aug 2025 13:15:04 +0800 Subject: [PATCH 29/43] style: format code --- src/types/tdigest.h | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/types/tdigest.h b/src/types/tdigest.h index e7c637ce54f..b2368073b56 100644 --- a/src/types/tdigest.h +++ b/src/types/tdigest.h @@ -172,7 +172,7 @@ inline Status TDigestRank(TD&& td, const std::vector& inputs, std::vecto auto iter = td.End(); while (i > 0) { auto centroid = GET_OR_RET(iter->GetCentroid()); - + if (centroid.mean > inputs[indices[i - 1]]) { // mean > input, accumulate weight and move to prev centroid cumulative_weight += centroid.weight; @@ -181,7 +181,7 @@ inline Status TDigestRank(TD&& td, const std::vector& inputs, std::vecto cumulative_weight += centroid.weight; auto current_mean = centroid.mean; auto current_mean_cumulative_weight = cumulative_weight + centroid.weight / 2; - + // handle all the prev centroids which has the same mean while (!iter->IsAtBegin() && iter->Prev()) { auto next_centroid = GET_OR_RET(iter->GetCentroid()); @@ -208,7 +208,7 @@ inline Status TDigestRank(TD&& td, const std::vector& inputs, std::vecto result[indices[i - 1]] = static_cast(cumulative_weight); i--; } - + if (iter->IsAtBegin()) { break; } From 07836fd6428552f2684486360f1982d5157047c4 Mon Sep 17 00:00:00 2001 From: donghao526 Date: Fri, 19 Sep 2025 15:54:12 +0800 Subject: [PATCH 30/43] feat: sort the input using map in revrank --- src/types/tdigest.h | 77 ++++++++++++++++++++++----------------------- 1 file changed, 38 insertions(+), 39 deletions(-) diff --git a/src/types/tdigest.h b/src/types/tdigest.h index b2368073b56..e3f0c65ac45 100644 --- a/src/types/tdigest.h +++ b/src/types/tdigest.h @@ -22,6 +22,7 @@ #include +#include #include #include @@ -154,71 +155,69 @@ inline StatusOr TDigestQuantile(TD&& td, double q) { template inline Status TDigestRank(TD&& td, const std::vector& inputs, std::vector& result) { - std::vector indices(inputs.size()); - std::iota(indices.begin(), indices.end(), 0); - std::sort(indices.begin(), indices.end(), [&inputs](size_t a, size_t b) { return inputs[a] < inputs[b]; }); + std::map> value_to_indices; + for (size_t i = 0; i < inputs.size(); ++i) { + value_to_indices[inputs[i]].push_back(i); + } - result.resize(inputs.size()); - size_t i = indices.size(); double cumulative_weight = 0; + result.resize(inputs.size()); + auto it = value_to_indices.begin(); // handle inputs larger than maximum - while (i > 0 && inputs[indices[i - 1]] > td.Max()) { - result[indices[i - 1]] = -1; - i--; + while (it != value_to_indices.end() && it->first < td.Min()) { + for (auto index : it->second) { + result[index] = -1; + } + ++it; } - // reverse iterate through centroids and calculate reverse rank for each input - auto iter = td.End(); - while (i > 0) { - auto centroid = GET_OR_RET(iter->GetCentroid()); + auto iter = td.Begin(); - if (centroid.mean > inputs[indices[i - 1]]) { - // mean > input, accumulate weight and move to prev centroid - cumulative_weight += centroid.weight; - } else if (centroid.mean == inputs[indices[i - 1]]) { - // mean == input, calculate reverse rank with half weight of current centroid - cumulative_weight += centroid.weight; + while (iter->Valid() && it != value_to_indices.end()) { + auto centroid = GET_OR_RET(iter->GetCentroid()); + auto input_value = it->first; + if (centroid.mean == input_value) { auto current_mean = centroid.mean; auto current_mean_cumulative_weight = cumulative_weight + centroid.weight / 2; + cumulative_weight += centroid.weight; // handle all the prev centroids which has the same mean - while (!iter->IsAtBegin() && iter->Prev()) { + while (iter->Next()) { auto next_centroid = GET_OR_RET(iter->GetCentroid()); if (current_mean != next_centroid.mean) { // move back to the last equal centroid, because we will process it in the next loop - iter->Next(); + iter->Prev(); break; } current_mean_cumulative_weight += centroid.weight / 2; cumulative_weight += centroid.weight; } - // assign the reverse rank for the inputs[indices[i - 1]] - result[indices[i - 1]] = static_cast(current_mean_cumulative_weight); - i--; - // handle the prev inputs which has the same value - while ((i > 0) && (inputs[indices[i]] == inputs[indices[i - 1]])) { - result[indices[i - 1]] = result[indices[i]]; - i--; + for (auto index : it->second) { + result[index] = static_cast(current_mean_cumulative_weight); } + ++it; + iter->Next(); + } else if (centroid.mean < input_value) { + cumulative_weight += centroid.weight; + iter->Next(); } else { - // mean < input, calculate reverse rank - result[indices[i - 1]] = static_cast(cumulative_weight); - i--; - } - - if (iter->IsAtBegin()) { - break; + for (auto index : it->second) { + result[index] = static_cast(cumulative_weight); + } + ++it; } - iter->Prev(); } // handle inputs less than minimum - while (i > 0) { - result[indices[i - 1]] = static_cast(td.TotalWeight()); - i--; + while (it != value_to_indices.end()) { + for (auto index : it->second) { + result[index] = -1; + } + ++it; } + return Status::OK(); -} \ No newline at end of file +} From f4a9c537d1758237ab7b825b160bd3ffb23de4ef Mon Sep 17 00:00:00 2001 From: donghao526 Date: Sat, 25 Oct 2025 21:21:24 +0800 Subject: [PATCH 31/43] feat: add the support of TDIGEST.REVRANK command --- src/types/redis_tdigest.cc | 10 ++++------ src/types/tdigest.h | 25 ++++++++++++------------- 2 files changed, 16 insertions(+), 19 deletions(-) diff --git a/src/types/redis_tdigest.cc b/src/types/redis_tdigest.cc index feded24fcae..81bf36caaca 100644 --- a/src/types/redis_tdigest.cc +++ b/src/types/redis_tdigest.cc @@ -67,20 +67,18 @@ class DummyCentroids { if (Valid()) { std::advance(iter_, 1); } - return iter_ != centroids_.cend(); + return Valid(); } // The Prev function can only be called for item is not cend, // because we must guarantee the iterator to be inside the valid range before iteration. bool Prev() { - if (Valid() && iter_ != centroids_.cbegin()) { + if (Valid()) { std::advance(iter_, -1); } return Valid(); } - bool Valid() const { return iter_ != centroids_.cend(); } - bool IsAtBegin() const { return iter_ == centroids_.cbegin(); } - + bool Valid() const { return iter_ < centroids_.cend() && iter_ >= centroids_.cbegin(); } StatusOr GetCentroid() const { if (iter_ == centroids_.cend()) { return {::Status::NotOK, "invalid iterator during decoding tdigest centroid"}; @@ -244,7 +242,7 @@ rocksdb::Status TDigest::RevRank(engine::Context& ctx, const Slice& digest_name, } auto dump_centroids = DummyCentroids(metadata, centroids); - auto status = TDigestRank(dump_centroids, inputs, result); + auto status = TDigestRevRank(dump_centroids, inputs, result); if (!status) { return rocksdb::Status::InvalidArgument(status.Msg()); } diff --git a/src/types/tdigest.h b/src/types/tdigest.h index e3f0c65ac45..8de847cbbfe 100644 --- a/src/types/tdigest.h +++ b/src/types/tdigest.h @@ -154,7 +154,7 @@ inline StatusOr TDigestQuantile(TD&& td, double q) { } template -inline Status TDigestRank(TD&& td, const std::vector& inputs, std::vector& result) { +inline Status TDigestRevRank(TD&& td, const std::vector& inputs, std::vector& result) { std::map> value_to_indices; for (size_t i = 0; i < inputs.size(); ++i) { value_to_indices[inputs[i]].push_back(i); @@ -162,19 +162,18 @@ inline Status TDigestRank(TD&& td, const std::vector& inputs, std::vecto double cumulative_weight = 0; result.resize(inputs.size()); - auto it = value_to_indices.begin(); + auto it = value_to_indices.rbegin(); // handle inputs larger than maximum - while (it != value_to_indices.end() && it->first < td.Min()) { + while (it != value_to_indices.rend() && it->first > td.Max()) { for (auto index : it->second) { result[index] = -1; } ++it; } - auto iter = td.Begin(); - - while (iter->Valid() && it != value_to_indices.end()) { + auto iter = td.End(); + while (iter->Valid() && it != value_to_indices.rend()) { auto centroid = GET_OR_RET(iter->GetCentroid()); auto input_value = it->first; if (centroid.mean == input_value) { @@ -183,11 +182,11 @@ inline Status TDigestRank(TD&& td, const std::vector& inputs, std::vecto cumulative_weight += centroid.weight; // handle all the prev centroids which has the same mean - while (iter->Next()) { + while (iter->Prev()) { auto next_centroid = GET_OR_RET(iter->GetCentroid()); if (current_mean != next_centroid.mean) { // move back to the last equal centroid, because we will process it in the next loop - iter->Prev(); + iter->Next(); break; } current_mean_cumulative_weight += centroid.weight / 2; @@ -199,10 +198,10 @@ inline Status TDigestRank(TD&& td, const std::vector& inputs, std::vecto result[index] = static_cast(current_mean_cumulative_weight); } ++it; - iter->Next(); - } else if (centroid.mean < input_value) { + iter->Prev(); + } else if (centroid.mean > input_value) { cumulative_weight += centroid.weight; - iter->Next(); + iter->Prev(); } else { for (auto index : it->second) { result[index] = static_cast(cumulative_weight); @@ -212,9 +211,9 @@ inline Status TDigestRank(TD&& td, const std::vector& inputs, std::vecto } // handle inputs less than minimum - while (it != value_to_indices.end()) { + while (it != value_to_indices.rend()) { for (auto index : it->second) { - result[index] = -1; + result[index] = static_cast(td.TotalWeight()); } ++it; } From e3629d969b9ad70be3747a143a3f63e568c8bcbb Mon Sep 17 00:00:00 2001 From: donghao526 Date: Sat, 25 Oct 2025 22:09:19 +0800 Subject: [PATCH 32/43] feat: add the support of TDIGEST.REVRANK command --- src/commands/cmd_tdigest.cc | 2 +- src/types/tdigest.h | 22 ++++++++++------------ 2 files changed, 11 insertions(+), 13 deletions(-) diff --git a/src/commands/cmd_tdigest.cc b/src/commands/cmd_tdigest.cc index 4bc1f2df9c8..ab90559e07e 100644 --- a/src/commands/cmd_tdigest.cc +++ b/src/commands/cmd_tdigest.cc @@ -201,7 +201,7 @@ class CommandTDigestRevRank : public Commander { return {Status::RedisExecErr, s.ToString()}; } - if (result.data()) { + if (!result.empty()) { std::vector rev_ranks; rev_ranks.reserve(result.size()); for (const auto v : result) { diff --git a/src/types/tdigest.h b/src/types/tdigest.h index 8de847cbbfe..37bc19f7854 100644 --- a/src/types/tdigest.h +++ b/src/types/tdigest.h @@ -153,6 +153,12 @@ inline StatusOr TDigestQuantile(TD&& td, double q) { return Lerp(lc.mean, rc.mean, diff); } +inline void assignRankForEqualInputs(const std::vector& indices, double cumulative_weight, std::vector& result) { + for (auto index : indices) { + result[index] = static_cast(cumulative_weight); + } +} + template inline Status TDigestRevRank(TD&& td, const std::vector& inputs, std::vector& result) { std::map> value_to_indices; @@ -166,9 +172,7 @@ inline Status TDigestRevRank(TD&& td, const std::vector& inputs, std::ve // handle inputs larger than maximum while (it != value_to_indices.rend() && it->first > td.Max()) { - for (auto index : it->second) { - result[index] = -1; - } + assignRankForEqualInputs(it->second, -1, result); ++it; } @@ -194,27 +198,21 @@ inline Status TDigestRevRank(TD&& td, const std::vector& inputs, std::ve } // handle the prev inputs which has the same value - for (auto index : it->second) { - result[index] = static_cast(current_mean_cumulative_weight); - } + assignRankForEqualInputs(it->second, current_mean_cumulative_weight, result); ++it; iter->Prev(); } else if (centroid.mean > input_value) { cumulative_weight += centroid.weight; iter->Prev(); } else { - for (auto index : it->second) { - result[index] = static_cast(cumulative_weight); - } + assignRankForEqualInputs(it->second, cumulative_weight, result); ++it; } } // handle inputs less than minimum while (it != value_to_indices.rend()) { - for (auto index : it->second) { - result[index] = static_cast(td.TotalWeight()); - } + assignRankForEqualInputs(it->second, td.TotalWeight(), result); ++it; } From ae0562306ed60d5d668f7837a65d8712291ed595 Mon Sep 17 00:00:00 2001 From: donghao526 Date: Sat, 25 Oct 2025 23:36:28 +0800 Subject: [PATCH 33/43] fix: fix format --- src/types/tdigest.h | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/types/tdigest.h b/src/types/tdigest.h index 37bc19f7854..4e46cc3ffd3 100644 --- a/src/types/tdigest.h +++ b/src/types/tdigest.h @@ -153,7 +153,8 @@ inline StatusOr TDigestQuantile(TD&& td, double q) { return Lerp(lc.mean, rc.mean, diff); } -inline void assignRankForEqualInputs(const std::vector& indices, double cumulative_weight, std::vector& result) { +inline void assignRankForEqualInputs(const std::vector& indices, double cumulative_weight, + std::vector& result) { for (auto index : indices) { result[index] = static_cast(cumulative_weight); } From 0cf8c8a7c354f26b878d1f9bfe6857dc187f295d Mon Sep 17 00:00:00 2001 From: donghao526 Date: Sun, 26 Oct 2025 10:46:35 +0800 Subject: [PATCH 34/43] fix: fix clang-tidy --- src/types/tdigest.h | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/src/types/tdigest.h b/src/types/tdigest.h index 4e46cc3ffd3..9557e0bb492 100644 --- a/src/types/tdigest.h +++ b/src/types/tdigest.h @@ -153,7 +153,7 @@ inline StatusOr TDigestQuantile(TD&& td, double q) { return Lerp(lc.mean, rc.mean, diff); } -inline void assignRankForEqualInputs(const std::vector& indices, double cumulative_weight, +inline void AssignRankForEqualInputs(const std::vector& indices, double cumulative_weight, std::vector& result) { for (auto index : indices) { result[index] = static_cast(cumulative_weight); @@ -173,7 +173,7 @@ inline Status TDigestRevRank(TD&& td, const std::vector& inputs, std::ve // handle inputs larger than maximum while (it != value_to_indices.rend() && it->first > td.Max()) { - assignRankForEqualInputs(it->second, -1, result); + AssignRankForEqualInputs(it->second, -1, result); ++it; } @@ -199,21 +199,21 @@ inline Status TDigestRevRank(TD&& td, const std::vector& inputs, std::ve } // handle the prev inputs which has the same value - assignRankForEqualInputs(it->second, current_mean_cumulative_weight, result); + AssignRankForEqualInputs(it->second, current_mean_cumulative_weight, result); ++it; iter->Prev(); } else if (centroid.mean > input_value) { cumulative_weight += centroid.weight; iter->Prev(); } else { - assignRankForEqualInputs(it->second, cumulative_weight, result); + AssignRankForEqualInputs(it->second, cumulative_weight, result); ++it; } } // handle inputs less than minimum while (it != value_to_indices.rend()) { - assignRankForEqualInputs(it->second, td.TotalWeight(), result); + AssignRankForEqualInputs(it->second, td.TotalWeight(), result); ++it; } From c6d5fc7dd15bbeb5d96643c8ebd2b01f8043f0e7 Mon Sep 17 00:00:00 2001 From: donghao526 Date: Mon, 27 Oct 2025 20:58:21 +0800 Subject: [PATCH 35/43] feat: add the support of TDIGEST.REVRANK command --- src/types/redis_tdigest.cc | 10 +++++++--- src/types/tdigest.h | 12 +++++++++--- 2 files changed, 16 insertions(+), 6 deletions(-) diff --git a/src/types/redis_tdigest.cc b/src/types/redis_tdigest.cc index 81bf36caaca..992775ec51f 100644 --- a/src/types/redis_tdigest.cc +++ b/src/types/redis_tdigest.cc @@ -67,18 +67,22 @@ class DummyCentroids { if (Valid()) { std::advance(iter_, 1); } - return Valid(); + return iter_ != centroids_.cend();; + } + + bool IsBegin() { + return iter_ == centroids_.cbegin(); } // The Prev function can only be called for item is not cend, // because we must guarantee the iterator to be inside the valid range before iteration. bool Prev() { - if (Valid()) { + if (Valid() && iter_ != centroids_.cbegin()) { std::advance(iter_, -1); } return Valid(); } - bool Valid() const { return iter_ < centroids_.cend() && iter_ >= centroids_.cbegin(); } + bool Valid() const { return iter_ != centroids_.cend(); } StatusOr GetCentroid() const { if (iter_ == centroids_.cend()) { return {::Status::NotOK, "invalid iterator during decoding tdigest centroid"}; diff --git a/src/types/tdigest.h b/src/types/tdigest.h index 9557e0bb492..33f697d0324 100644 --- a/src/types/tdigest.h +++ b/src/types/tdigest.h @@ -187,23 +187,29 @@ inline Status TDigestRevRank(TD&& td, const std::vector& inputs, std::ve cumulative_weight += centroid.weight; // handle all the prev centroids which has the same mean - while (iter->Prev()) { + while (!iter->IsBegin() && iter->Prev()) { auto next_centroid = GET_OR_RET(iter->GetCentroid()); if (current_mean != next_centroid.mean) { // move back to the last equal centroid, because we will process it in the next loop iter->Next(); break; } - current_mean_cumulative_weight += centroid.weight / 2; - cumulative_weight += centroid.weight; + current_mean_cumulative_weight += next_centroid.weight / 2; + cumulative_weight += next_centroid.weight; } // handle the prev inputs which has the same value AssignRankForEqualInputs(it->second, current_mean_cumulative_weight, result); ++it; + if (iter->IsBegin()) { + break; + } iter->Prev(); } else if (centroid.mean > input_value) { cumulative_weight += centroid.weight; + if (iter->IsBegin()) { + break; + } iter->Prev(); } else { AssignRankForEqualInputs(it->second, cumulative_weight, result); From 62dfbb53cd7d49dae573cfe16f75fddc7453c0ff Mon Sep 17 00:00:00 2001 From: donghao526 Date: Mon, 27 Oct 2025 22:11:59 +0800 Subject: [PATCH 36/43] feat: add the support of TDIGEST.REVRANK command --- src/types/redis_tdigest.cc | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/src/types/redis_tdigest.cc b/src/types/redis_tdigest.cc index 992775ec51f..2d45a58d853 100644 --- a/src/types/redis_tdigest.cc +++ b/src/types/redis_tdigest.cc @@ -67,12 +67,10 @@ class DummyCentroids { if (Valid()) { std::advance(iter_, 1); } - return iter_ != centroids_.cend();; + return iter_ != centroids_.cend(); } - bool IsBegin() { - return iter_ == centroids_.cbegin(); - } + bool IsBegin() { return iter_ == centroids_.cbegin(); } // The Prev function can only be called for item is not cend, // because we must guarantee the iterator to be inside the valid range before iteration. From 3471dac9adc72de9b524ff7ceb745222714c368e Mon Sep 17 00:00:00 2001 From: donghao526 Date: Tue, 28 Oct 2025 10:32:22 +0800 Subject: [PATCH 37/43] fix: remove unnecessary empty check for tdigest.revrank --- src/commands/cmd_tdigest.cc | 14 +++++--------- 1 file changed, 5 insertions(+), 9 deletions(-) diff --git a/src/commands/cmd_tdigest.cc b/src/commands/cmd_tdigest.cc index ab90559e07e..4cdadc71580 100644 --- a/src/commands/cmd_tdigest.cc +++ b/src/commands/cmd_tdigest.cc @@ -201,16 +201,12 @@ class CommandTDigestRevRank : public Commander { return {Status::RedisExecErr, s.ToString()}; } - if (!result.empty()) { - std::vector rev_ranks; - rev_ranks.reserve(result.size()); - for (const auto v : result) { - rev_ranks.push_back(redis::Integer(v)); - } - *output = redis::Array(rev_ranks); - } else { - *output = redis::BulkString("nan"); + std::vector rev_ranks; + rev_ranks.reserve(result.size()); + for (const auto v : result) { + rev_ranks.push_back(redis::Integer(v)); } + *output = redis::Array(rev_ranks); return Status::OK(); } From e943ab8a27052d131669dcfe5eb5380483354adf Mon Sep 17 00:00:00 2001 From: donghao526 Date: Wed, 29 Oct 2025 10:01:23 +0800 Subject: [PATCH 38/43] feat: use a stable way to compare two double values --- src/types/tdigest.h | 23 ++++++++++++++++++++--- 1 file changed, 20 insertions(+), 3 deletions(-) diff --git a/src/types/tdigest.h b/src/types/tdigest.h index 33f697d0324..16b5adcfa4d 100644 --- a/src/types/tdigest.h +++ b/src/types/tdigest.h @@ -160,9 +160,26 @@ inline void AssignRankForEqualInputs(const std::vector& indices, double } } +inline int DoubleCompare(double a, double b, double rel_eps = 1e-12, double abs_eps = 1e-9) { + double diff = a - b; + double adiff = std::abs(diff); + if (adiff <= abs_eps) return 0; + double maxab = std::max(std::abs(a), std::abs(b)); + if (adiff <= maxab * rel_eps) return 0; + return (diff < 0) ? -1 : 1; +} + +inline bool DoubleEqual(double a, double b, double rel_eps = 1e-12, double abs_eps = 1e-9) { + return DoubleCompare(a, b, rel_eps, abs_eps) == 0; +} + +struct DoubleComparator { + bool operator()(const double& a, const double& b) const { return DoubleCompare(a, b) == -1; } +}; + template inline Status TDigestRevRank(TD&& td, const std::vector& inputs, std::vector& result) { - std::map> value_to_indices; + std::map, DoubleComparator> value_to_indices; for (size_t i = 0; i < inputs.size(); ++i) { value_to_indices[inputs[i]].push_back(i); } @@ -181,7 +198,7 @@ inline Status TDigestRevRank(TD&& td, const std::vector& inputs, std::ve while (iter->Valid() && it != value_to_indices.rend()) { auto centroid = GET_OR_RET(iter->GetCentroid()); auto input_value = it->first; - if (centroid.mean == input_value) { + if (DoubleEqual(centroid.mean, input_value)) { auto current_mean = centroid.mean; auto current_mean_cumulative_weight = cumulative_weight + centroid.weight / 2; cumulative_weight += centroid.weight; @@ -205,7 +222,7 @@ inline Status TDigestRevRank(TD&& td, const std::vector& inputs, std::ve break; } iter->Prev(); - } else if (centroid.mean > input_value) { + } else if (DoubleCompare(centroid.mean, input_value) > 0) { cumulative_weight += centroid.weight; if (iter->IsBegin()) { break; From 85e4c4426ffd91d95624b638e63600ccbb9f23d6 Mon Sep 17 00:00:00 2001 From: Hao Dong Date: Thu, 30 Oct 2025 09:26:45 +0800 Subject: [PATCH 39/43] fix: fix the comments for clarity Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- tests/gocase/unit/type/tdigest/tdigest_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/gocase/unit/type/tdigest/tdigest_test.go b/tests/gocase/unit/type/tdigest/tdigest_test.go index 91c453bb33f..aa0462d53ee 100644 --- a/tests/gocase/unit/type/tdigest/tdigest_test.go +++ b/tests/gocase/unit/type/tdigest/tdigest_test.go @@ -544,7 +544,7 @@ func tdigestTests(t *testing.T, configs util.KvrocksServerConfigs) { require.EqualValues(t, rank, expected[i]) } - // Test with set_contains several identical elements + // Test with set containing several identical elements require.NoError(t, rdb.Do(ctx, "TDIGEST.ADD", key, "10", "10", "10", "20", "20").Err()) rsp = rdb.Do(ctx, "TDIGEST.REVRANK", key, "10", "20") require.NoError(t, rsp.Err()) From ac320809e2c2665583023e18af28af0690a61183 Mon Sep 17 00:00:00 2001 From: Hao Dong Date: Thu, 30 Oct 2025 09:27:43 +0800 Subject: [PATCH 40/43] fix: fix the comments for clarity Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- tests/gocase/unit/type/tdigest/tdigest_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/gocase/unit/type/tdigest/tdigest_test.go b/tests/gocase/unit/type/tdigest/tdigest_test.go index aa0462d53ee..129be055405 100644 --- a/tests/gocase/unit/type/tdigest/tdigest_test.go +++ b/tests/gocase/unit/type/tdigest/tdigest_test.go @@ -571,7 +571,7 @@ func tdigestTests(t *testing.T, configs util.KvrocksServerConfigs) { require.EqualValues(t, rank, expected[i]) } - // Test with set_contains different elements + // Test with set containing different elements key2 := keyPrefix + "test2" require.NoError(t, rdb.Do(ctx, "TDIGEST.CREATE", key2, "compression", "100").Err()) require.NoError(t, rdb.Do(ctx, "TDIGEST.ADD", key2, "10", "20", "30", "40", "50", "60").Err()) From 79d2e3330435578310882eca7897e8dc03f3a3a4 Mon Sep 17 00:00:00 2001 From: Hao Dong Date: Thu, 30 Oct 2025 09:30:13 +0800 Subject: [PATCH 41/43] fix: use DoubleEqual to ensure the two centorid.mean comparison stable Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- src/types/tdigest.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/types/tdigest.h b/src/types/tdigest.h index 16b5adcfa4d..84260d3d85b 100644 --- a/src/types/tdigest.h +++ b/src/types/tdigest.h @@ -206,7 +206,7 @@ inline Status TDigestRevRank(TD&& td, const std::vector& inputs, std::ve // handle all the prev centroids which has the same mean while (!iter->IsBegin() && iter->Prev()) { auto next_centroid = GET_OR_RET(iter->GetCentroid()); - if (current_mean != next_centroid.mean) { + if (!DoubleEqual(current_mean, next_centroid.mean)) { // move back to the last equal centroid, because we will process it in the next loop iter->Next(); break; From 625510d1123ac21b842df89c5ace0503ad7de22a Mon Sep 17 00:00:00 2001 From: Hao Dong Date: Thu, 30 Oct 2025 09:31:06 +0800 Subject: [PATCH 42/43] fix: Corrected phrasing for proper naming convention Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- tests/cppunit/types/tdigest_test.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/cppunit/types/tdigest_test.cc b/tests/cppunit/types/tdigest_test.cc index 09b27fd70cd..644693b89d3 100644 --- a/tests/cppunit/types/tdigest_test.cc +++ b/tests/cppunit/types/tdigest_test.cc @@ -322,7 +322,7 @@ TEST_F(RedisTDigestTest, RevRank_on_the_set_contains_different_elements) { ASSERT_TRUE(status.ok()) << status.ToString(); } -TEST_F(RedisTDigestTest, RevRank_on_the_set_contains_several_identical_elements) { +TEST_F(RedisTDigestTest, RevRank_on_the_set_containing_several_identical_elements) { std::string test_digest_name = "test_digest_revrank" + std::to_string(util::GetTimeStampMS()); bool exists = false; auto status = tdigest_->Create(*ctx_, test_digest_name, {100}, &exists); From cc013737924cb2faa0ab8dc137e23332822e53af Mon Sep 17 00:00:00 2001 From: Hao Dong Date: Thu, 30 Oct 2025 09:31:23 +0800 Subject: [PATCH 43/43] fix: Corrected phrasing for proper naming convention Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- tests/cppunit/types/tdigest_test.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/cppunit/types/tdigest_test.cc b/tests/cppunit/types/tdigest_test.cc index 644693b89d3..ae4bf29e9e3 100644 --- a/tests/cppunit/types/tdigest_test.cc +++ b/tests/cppunit/types/tdigest_test.cc @@ -299,7 +299,7 @@ TEST_F(RedisTDigestTest, Quantile_returns_nan_on_empty_tdigest) { ASSERT_FALSE(result.quantiles) << "should not have quantiles with empty tdigest"; } -TEST_F(RedisTDigestTest, RevRank_on_the_set_contains_different_elements) { +TEST_F(RedisTDigestTest, RevRank_on_the_set_containing_different_elements) { std::string test_digest_name = "test_digest_revrank" + std::to_string(util::GetTimeStampMS()); bool exists = false; auto status = tdigest_->Create(*ctx_, test_digest_name, {100}, &exists);