From 9de9d1661c13821cde141417a7f8aa8ddf624b2b Mon Sep 17 00:00:00 2001 From: Zach McCoy Date: Mon, 31 Aug 2026 11:56:57 -0500 Subject: [PATCH 1/6] Merge and tests for merge with questions, cherry picked from old rebase branch Signed-off-by: Zach McCoy --- src/cms/cms_command_handler.rs | 94 +++++++++++++++++++++++++++++++++- src/cms/utils.rs | 34 +++++++++++- src/lib.rs | 6 +++ tests/test_cms_basic.py | 22 +++++++- tests/test_cms_command.py | 9 ++++ 5 files changed, 162 insertions(+), 3 deletions(-) diff --git a/src/cms/cms_command_handler.rs b/src/cms/cms_command_handler.rs index 077b88f2..37fbeaf6 100644 --- a/src/cms/cms_command_handler.rs +++ b/src/cms/cms_command_handler.rs @@ -1,5 +1,6 @@ use valkey_module::{ - Context, NotifyEvent, ValkeyError, ValkeyResult, ValkeyString, ValkeyValue, VALKEY_OK, + key::ValkeyKey, Context, NotifyEvent, ValkeyError, ValkeyResult, ValkeyString, ValkeyValue, + VALKEY_OK, }; use crate::cms::data_type::CMS_TYPE; @@ -13,6 +14,7 @@ enum Replications { enum Operation { Initialization { replications: Replications }, Increment, + Merge, } fn replicate_and_notify_events(ctx: &Context, key_name: &ValkeyString, operation: Operation) { @@ -52,6 +54,11 @@ fn replicate_and_notify_events(ctx: &Context, key_name: &ValkeyString, operation ctx.replicate_verbatim(); ctx.notify_keyspace_event(NotifyEvent::GENERIC, utils::INCR_EVENT, key_name); } + Operation::Merge => { + //TODO::How should this replication be done, and why vs verbatim + ctx.replicate_verbatim(); + ctx.notify_keyspace_event(NotifyEvent::GENERIC, utils::MERGE_EVENT, key_name); + } } } @@ -238,3 +245,88 @@ pub fn cms_query(ctx: &Context, args: Vec) -> ValkeyResult { } } } + +/// Function that implements logic to handle the CMS.MERGE command. +pub fn cms_merge(ctx: &Context, args: Vec) -> ValkeyResult { + let args_count = args.len(); + if args_count < 5 { + return Err(ValkeyError::WrongArity); + } + + //This must already be initialized. + let destination_key = &args[1]; + + let number_of_keys_value = args[2].to_string_lossy().parse::()?; + + //Indexes 3 -> 3 + N-1 are keys to merge + let sketch_end_index = 3 + number_of_keys_value - 1; + //Up to non inclusive grab 3 up to sketch_end_index + let source_keys: Vec<&ValkeyString> = args[3..=sketch_end_index].iter().collect(); + + //Then Parse the optional WEIGHTS section of the command. WEIGHTS is at sketch_end + 1 + let passed_in_weights = if sketch_end_index + 1 == args_count { + Vec::new() + } else { + //Make sure we have at least WEIGHT weight left + let weight_args_left = args_count - sketch_end_index - 1; + if weight_args_left < 2 { + return Err(ValkeyError::WrongArity); + } + + //There should be at least 2 args left WEIGHT weight [weight ...] + let weights_keyword_index = sketch_end_index + 1; + let weights_keyword = args[weights_keyword_index].to_string_lossy(); + if weights_keyword.to_uppercase() != "WEIGHTS" { + return Err(ValkeyError::Str("ERR invalid argument")); + } + let weights_start = weights_keyword_index + 1; + let weights_args = args[weights_start..].iter(); + + let weights: Vec = weights_args + .map(|weight| { + weight + .to_string_lossy() + .parse::() + .map_err(|_| ValkeyError::Str("ERR invalid weight value")) + }) + .collect::, _>>()?; + weights + }; + + let source_key_handles: Vec = + source_keys.iter().map(|key| ctx.open_key(key)).collect(); + let sketches_result: Result, ValkeyError> = source_key_handles + .iter() + .map(|key_handle| { + key_handle + .get_value::(&CMS_TYPE) + .and_then(|opt| opt.ok_or_else(|| ValkeyError::Str("ERR key does not exist"))) + }) + .collect(); + + let sketches: Vec<&CMSObject> = sketches_result?; + let destination_sketch = ctx + .open_key_writable(destination_key) + .get_value::(&CMS_TYPE) + .and_then(|opt| opt.ok_or_else(|| ValkeyError::Str("ERR key does not exist")))?; + + let sketches_with_weights: Vec<(&CMSObject, f64)> = sketches + .into_iter() + .enumerate() + .map(|(i, sketch)| { + let weight = passed_in_weights.get(i).copied().unwrap_or(1.0); + (sketch, weight) + }) + .collect(); + + //Mutates the destination_sketch's internal CMS to be the merge of the sketches_with_weights + //Impl note: We do not handle the weights yet in the called function, as the lib-source needs to change. + destination_sketch + .merge(&sketches_with_weights) + .map_err(|_| { + ValkeyError::Str("ERR destination key is not of the same width and/or depth") + })?; + + replicate_and_notify_events(ctx, destination_key, Operation::Merge); + VALKEY_OK +} diff --git a/src/cms/utils.rs b/src/cms/utils.rs index 240c10d7..341ce623 100644 --- a/src/cms/utils.rs +++ b/src/cms/utils.rs @@ -1,4 +1,4 @@ -use flowstats::CountMinSketch; +use flowstats::{traits::Sketch, CountMinSketch}; /// Client Errors pub const ERROR: &str = "ERROR"; @@ -12,11 +12,13 @@ pub const BAD_PROBABILITY: &str = "ERR bad probability"; pub const PROBABILITY_RANGE: &str = "ERR probability rate should be between 0 and 1"; pub const KEY_EXISTS: &str = "ERR Target key name already exists."; pub const BAD_INCREMENT: &str = "ERR bad increment"; +pub const MERGE_FAILURE: &str = "ERR width / depth is not equal across sketches"; ///Keyspace Notification Events pub const INITBYPROB_EVENT: &str = "countminsketch.initbyprob"; pub const INITBYDIM_EVENT: &str = "countminsketch.initbydim"; pub const INCR_EVENT: &str = "countminsketch.incrby"; +pub const MERGE_EVENT: &str = "countminsketch.merge"; #[derive(Debug, PartialEq)] pub enum CMSError { @@ -24,6 +26,7 @@ pub enum CMSError { InvalidDepth, InvalidErrorRate, InvalidProbability, + MergeFailed, } impl CMSError { @@ -33,6 +36,7 @@ impl CMSError { CMSError::InvalidDepth => BAD_DEPTH, CMSError::InvalidErrorRate => ERROR_RATE_RANGE, CMSError::InvalidProbability => PROBABILITY_RANGE, + CMSError::MergeFailed => MERGE_FAILURE, } } } @@ -103,6 +107,34 @@ impl CMSObject { pub fn estimate(&self, item: &[u8]) -> u64 { self.cms.estimate_item(item) } + + pub fn merge(&mut self, sketches_and_weights: &[(&CMSObject, f64)]) -> Result<(), CMSError> { + //Pre-check this so dest is the same size as the others. + if sketches_and_weights + .iter() + .any(|sketch| sketch.0.width() != self.width() || sketch.0.depth != self.depth()) + { + return Err(CMSError::MergeFailed); + } + + let (head, tail) = sketches_and_weights + .split_first() + .ok_or(CMSError::MergeFailed)?; + + let mut merged_sketch = head.0.cms.sketch.clone(); + + for (cms, _weight) in tail.iter() { + merged_sketch + .merge(&cms.cms.sketch) + .map_err(|_| CMSError::MergeFailed)?; + } + + self.cms = CMS { + sketch: merged_sketch, + }; + + Ok(()) + } } struct CMS { diff --git a/src/lib.rs b/src/lib.rs index 8fbd05fa..93233467 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -114,6 +114,11 @@ fn cms_query_command(ctx: &Context, args: Vec) -> ValkeyResult { cms_command_handler::cms_query(ctx, args) } +/// Command handler for CMS.MERGE [ ...] [WEIGHTS [ ...]] +fn cms_merge_command(ctx: &Context, args: Vec) -> ValkeyResult { + cms_command_handler::cms_merge(ctx, args) +} + /// /// Module Info /// @@ -152,6 +157,7 @@ valkey_module! { ["CMS.INITBYPROB", cms_initbyprob_command, "write fast deny-oom", 1, 1, 1, "fast write cms"], ["CMS.INCRBY", cms_incrby_command, "write fast deny-oom", 1, 1, 1, "write cms"], ["CMS.QUERY", cms_query_command, "readonly fast", 1, 1, 1, "read cms"], + ["CMS.MERGE", cms_merge_command, "write deny-oom", 1, 1, 1, "write cms"], ], configurations: [ i64: [ diff --git a/tests/test_cms_basic.py b/tests/test_cms_basic.py index 25c90d7b..51e9f5bd 100644 --- a/tests/test_cms_basic.py +++ b/tests/test_cms_basic.py @@ -32,7 +32,6 @@ def test_basic_dim(self): assert client.execute_command('CMS.QUERY sketch1 item1')[0] >= 1 - def test_basic_prob(self): client = self.server.get_new_client() module_loaded = False @@ -54,6 +53,27 @@ def test_basic_prob(self): #CMS guarantees that we have the frequency at LEAST the size of the increment for the item assert client.execute_command('CMS.QUERY sketch1 item1')[0] >= 1 + + def test_merge(self): + client = self.server.get_new_client() + module_loaded = False + module_list_data = client.execute_command('MODULE LIST') + for module in module_list_data: + if (module[b'name'] == b'bf'): + module_loaded = True + break + assert(module_loaded) + #Create the destination and sketches to be merged into the destination key + assert client.execute_command('CMS.INITBYDIM dest 10 5') == b'OK' + assert client.execute_command('CMS.INITBYDIM s1 10 5') == b'OK' + assert client.execute_command('CMS.INITBYDIM s2 10 5') == b'OK' + assert client.execute_command('CMS.INCRBY s1 a 1 b 2') == [1, 2] + assert client.execute_command('CMS.INCRBY s2 a 1 b 3') == [1, 3] + assert client.execute_command('CMS.MERGE dest 2 s1 s2') == b'OK' + assert client.execute_command('CMS.QUERY dest a') == [2] + assert client.execute_command('CMS.QUERY dest b') == [5] + + def test_module_data_type(self): # Validate the name of the Module data type. client = self.server.get_new_client() diff --git a/tests/test_cms_command.py b/tests/test_cms_command.py index 9c25c890..f311852b 100644 --- a/tests/test_cms_command.py +++ b/tests/test_cms_command.py @@ -32,6 +32,9 @@ def test_cms_command_error(self): ('CMS.INCRBY sketch item abc', 'bad increment'), ('CMS.INCRBY sketch item -1', 'bad increment'), + ('CMS.MERGE dest 2 s1 s2 WEIGHTSS 1', 'invalid argument'), + ('CMS.MERGE dest 2 s1 s2 WEIGHTS a', 'invalid weight value'), + # wrong number of arguments ('CMS.INITBYDIM', "wrong number of arguments for 'CMS.INITBYDIM' command"), ('CMS.INITBYDIM key', "wrong number of arguments for 'CMS.INITBYDIM' command"), @@ -50,6 +53,12 @@ def test_cms_command_error(self): ('CMS.QUERY', "wrong number of arguments for 'CMS.QUERY' command"), ('CMS.QUERY key', "wrong number of arguments for 'CMS.QUERY' command"), + + ('CMS.MERGE', 'wrong number of arguments for \'CMS.MERGE\' command'), + ('CMS.MERGE dest', 'wrong number of arguments for \'CMS.MERGE\' command'), + ('CMS.MERGE dest 2', 'wrong number of arguments for \'CMS.MERGE\' command'), + ('CMS.MERGE dest 2 s1', 'wrong number of arguments for \'CMS.MERGE\' command'), + ('CMS.MERGE dest 2 s1 s2 WEIGHTS', 'wrong number of arguments for \'CMS.MERGE\' command'), ] From 3249f681f79f8400f0b7233b22a2e0d22277edfa Mon Sep 17 00:00:00 2001 From: Zach McCoy Date: Mon, 31 Aug 2026 12:08:52 -0500 Subject: [PATCH 2/6] Use turbo shark to use ? on the same line as collect Signed-off-by: Zach McCoy --- src/cms/cms_command_handler.rs | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/cms/cms_command_handler.rs b/src/cms/cms_command_handler.rs index 37fbeaf6..fea7e845 100644 --- a/src/cms/cms_command_handler.rs +++ b/src/cms/cms_command_handler.rs @@ -295,16 +295,15 @@ pub fn cms_merge(ctx: &Context, args: Vec) -> ValkeyResult { let source_key_handles: Vec = source_keys.iter().map(|key| ctx.open_key(key)).collect(); - let sketches_result: Result, ValkeyError> = source_key_handles + let sketches = source_key_handles .iter() .map(|key_handle| { key_handle .get_value::(&CMS_TYPE) .and_then(|opt| opt.ok_or_else(|| ValkeyError::Str("ERR key does not exist"))) }) - .collect(); + .collect::, _>>()?; - let sketches: Vec<&CMSObject> = sketches_result?; let destination_sketch = ctx .open_key_writable(destination_key) .get_value::(&CMS_TYPE) From 0b0557b5e1fe6c87e403efc80f9a7e5e1341f84c Mon Sep 17 00:00:00 2001 From: Zach McCoy Date: Mon, 31 Aug 2026 12:12:44 -0500 Subject: [PATCH 3/6] Use double quotes in test like the others Signed-off-by: Zach McCoy --- tests/test_cms_command.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/tests/test_cms_command.py b/tests/test_cms_command.py index f311852b..e46fae82 100644 --- a/tests/test_cms_command.py +++ b/tests/test_cms_command.py @@ -54,11 +54,11 @@ def test_cms_command_error(self): ('CMS.QUERY', "wrong number of arguments for 'CMS.QUERY' command"), ('CMS.QUERY key', "wrong number of arguments for 'CMS.QUERY' command"), - ('CMS.MERGE', 'wrong number of arguments for \'CMS.MERGE\' command'), - ('CMS.MERGE dest', 'wrong number of arguments for \'CMS.MERGE\' command'), - ('CMS.MERGE dest 2', 'wrong number of arguments for \'CMS.MERGE\' command'), - ('CMS.MERGE dest 2 s1', 'wrong number of arguments for \'CMS.MERGE\' command'), - ('CMS.MERGE dest 2 s1 s2 WEIGHTS', 'wrong number of arguments for \'CMS.MERGE\' command'), + ('CMS.MERGE', "wrong number of arguments for 'CMS.MERGE' command"), + ('CMS.MERGE dest', "wrong number of arguments for 'CMS.MERGE' command"), + ('CMS.MERGE dest 2', "wrong number of arguments for 'CMS.MERGE' command"), + ('CMS.MERGE dest 2 s1', "wrong number of arguments for 'CMS.MERGE' command"), + ('CMS.MERGE dest 2 s1 s2 WEIGHTS', "wrong number of arguments for 'CMS.MERGE' command"), ] From 20f0ead6770a69eaba829db036a0c6dd12e20b4d Mon Sep 17 00:00:00 2001 From: Zach McCoy Date: Mon, 31 Aug 2026 12:30:58 -0500 Subject: [PATCH 4/6] Morph error string if parse fails, use float parse where possible Signed-off-by: Zach McCoy --- src/cms/cms_command_handler.rs | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/src/cms/cms_command_handler.rs b/src/cms/cms_command_handler.rs index fea7e845..87354660 100644 --- a/src/cms/cms_command_handler.rs +++ b/src/cms/cms_command_handler.rs @@ -256,7 +256,10 @@ pub fn cms_merge(ctx: &Context, args: Vec) -> ValkeyResult { //This must already be initialized. let destination_key = &args[1]; - let number_of_keys_value = args[2].to_string_lossy().parse::()?; + let number_of_keys_value = args[2] + .to_string_lossy() + .parse::() + .map_err(|_| ValkeyError::Str("ERR invalid number of keys value"))?; //Indexes 3 -> 3 + N-1 are keys to merge let sketch_end_index = 3 + number_of_keys_value - 1; @@ -285,8 +288,7 @@ pub fn cms_merge(ctx: &Context, args: Vec) -> ValkeyResult { let weights: Vec = weights_args .map(|weight| { weight - .to_string_lossy() - .parse::() + .parse_float() .map_err(|_| ValkeyError::Str("ERR invalid weight value")) }) .collect::, _>>()?; From 8e714883889f05c3c62adada9968aa4e647c2ce6 Mon Sep 17 00:00:00 2001 From: Zach McCoy Date: Mon, 31 Aug 2026 12:35:04 -0500 Subject: [PATCH 5/6] Add comment for CMS Object Signed-off-by: Zach McCoy --- src/cms/utils.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/cms/utils.rs b/src/cms/utils.rs index 341ce623..326e244f 100644 --- a/src/cms/utils.rs +++ b/src/cms/utils.rs @@ -108,6 +108,8 @@ impl CMSObject { self.cms.estimate_item(item) } + /// Merges the current CMS structure with a list of others as well as their weights. Weights are applied to the paired structure via multiplication first. + /// The values in each index of each of the K arrays internally then are added together into a single CMS structure. pub fn merge(&mut self, sketches_and_weights: &[(&CMSObject, f64)]) -> Result<(), CMSError> { //Pre-check this so dest is the same size as the others. if sketches_and_weights From a11fe1c731e7d1b65a74c1b3590e6420f9b32b31 Mon Sep 17 00:00:00 2001 From: Zach McCoy Date: Tue, 8 Sep 2026 19:38:00 -0500 Subject: [PATCH 6/6] Fix merge issue Signed-off-by: Zach McCoy --- src/cms/utils.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/cms/utils.rs b/src/cms/utils.rs index c63fbc25..2ae4612d 100644 --- a/src/cms/utils.rs +++ b/src/cms/utils.rs @@ -107,7 +107,7 @@ impl CMSObject { //Pre-check this so dest is the same size as the others. if sketches_and_weights .iter() - .any(|sketch| sketch.0.width() != self.width() || sketch.0.depth != self.depth()) + .any(|sketch| sketch.0.width != self.width || sketch.0.depth != self.depth) { return Err(CMSError::MergeFailed); }