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 <[email protected]>
Co-authored-by: dvora-h <[email protected]>
Co-authored-by: dvora-h <[email protected]>

Author: 3 people

redis/commands/bf/__init__.py

@@ -170,6 +170,7 @@ def __init__(self, client, **kwargs):
             TDIGEST_QUANTILE: parse_tdigest_quantile,
             TDIGEST_MIN: float,
             TDIGEST_MAX: float,
+            TDIGEST_TRIMMED_MEAN: float,
             TDIGEST_INFO: TDigestInfo,
         }
 

redis/commands/bf/commands.py

@@ -49,6 +49,7 @@
 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 @@ def info(self, key):
         """  # 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.

tests/test_bloom.py

@@ -388,6 +388,17 @@ def test_tdigest_cdf(client):
     assert 0.9 == round(client.tdigest().cdf("tDigest", 9.0), 1)
 
 
+@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")