Add TDIGEST.TRIMMED_MEAN (#2300)

* Add tdigest trimmed mean command with test

* Add skip version for test

* add to module callbacks

Co-authored-by: Alibi Shalgymbay <a.shalgymbay@mycar.kz>
Co-authored-by: dvora-h <67596500+dvora-h@users.noreply.github.com>
Co-authored-by: dvora-h <dvora.heller@redis.com>

3 files changed, 23 insertions, 0 deletions

diff --git a/redis/commands/bf/__init__.py b/redis/commands/bf/__init__.py
index 3169ba2..e41354b 100644
--- a/redis/commands/bf/__init__.py
+++ b/redis/commands/bf/__init__.py
@@ -170,6 +170,7 @@ class TDigestBloom(TDigestCommands, AbstractBloom):
             TDIGEST_QUANTILE: parse_tdigest_quantile,
             TDIGEST_MIN: float,
             TDIGEST_MAX: float,
+            TDIGEST_TRIMMED_MEAN: float,
             TDIGEST_INFO: TDigestInfo,
         }
 
diff --git a/redis/commands/bf/commands.py b/redis/commands/bf/commands.py
index 84a6b5f..0fe065d 100644
--- a/redis/commands/bf/commands.py
+++ b/redis/commands/bf/commands.py
@@ -49,6 +49,7 @@ TDIGEST_QUANTILE = "TDIGEST.QUANTILE"
 TDIGEST_MIN = "TDIGEST.MIN"
 TDIGEST_MAX = "TDIGEST.MAX"
 TDIGEST_INFO = "TDIGEST.INFO"
+TDIGEST_TRIMMED_MEAN = "TDIGEST.TRIMMED_MEAN"
 TDIGEST_MERGESTORE = "TDIGEST.MERGESTORE"
 
 
@@ -418,6 +419,16 @@ class TDigestCommands:
         """  # noqa
         return self.execute_command(TDIGEST_INFO, key)
 
+    def trimmed_mean(self, key, low_cut_quantile, high_cut_quantile):
+        """
+        Return mean value from the sketch, excluding observation values outside
+        the low and high cutoff quantiles.
+        For more information see `TDIGEST.TRIMMED_MEAN <https://redis.io/commands/tdigest.trimmed_mean>`_.
+        """  # noqa
+        return self.execute_command(
+            TDIGEST_TRIMMED_MEAN, key, low_cut_quantile, high_cut_quantile
+        )
+
     def mergestore(self, dest_key, numkeys, *sourcekeys, compression=False):
         """
         Merges all of the values from `sourcekeys` keys to `dest_key` sketch.
diff --git a/tests/test_bloom.py b/tests/test_bloom.py
index 1191542..12a4839 100644
--- a/tests/test_bloom.py
+++ b/tests/test_bloom.py
@@ -391,6 +391,17 @@ def test_tdigest_cdf(client):
 @pytest.mark.redismod
 @pytest.mark.experimental
 @skip_ifmodversion_lt("2.4.0", "bf")
+def test_tdigest_trimmed_mean(client):
+    assert client.tdigest().create("tDigest", 100)
+    # insert data-points into sketch
+    assert client.tdigest().add("tDigest", list(range(1, 10)), [1.0] * 10)
+    assert 5 == client.tdigest().trimmed_mean("tDigest", 0.1, 0.9)
+    assert 4.5 == client.tdigest().trimmed_mean("tDigest", 0.4, 0.5)
+
+
+@pytest.mark.redismod
+@pytest.mark.experimental
+@skip_ifmodversion_lt("2.4.0", "bf")
 def test_tdigest_mergestore(client):
     assert client.tdigest().create("sourcekey1", 100)
     assert client.tdigest().create("sourcekey2", 100)