Adding the Poisson distribution

prestodb · May 13, 2021 · 39c66e3 · 39c66e3
1 parent 77989bf
commit 39c66e3
Show file tree

Hide file tree

Showing 3 changed files with 63 additions and 0 deletions.
diff --git a/presto-docs/src/main/sphinx/functions/math.rst b/presto-docs/src/main/sphinx/functions/math.rst
@@ -92,6 +92,13 @@ Mathematical Functions
     a real value and the standard deviation must be a real and positive value.
     The probability p must lie on the interval (0, 1).
 
+.. function:: inverse_poisson_cdf(lambda, p) -> integer
+
+    Compute the inverse of the Poisson cdf with given lambda (mean) parameter for the cumulative
+    probability (p). It returns the value of n so that: P(N <= n; lambda) = p.
+    The lambda parameter must be a positive real number (of type DOUBLE).
+    The probability p must lie on the interval [0, 1).
+
 .. function:: normal_cdf(mean, sd, v) -> double
 
     Compute the Normal cdf with given mean and standard deviation (sd):  P(N < v; mean, sd).
@@ -130,6 +137,11 @@ Mathematical Functions
 
     Returns the constant Pi.
 
+.. function:: poisson_cdf(lambda, value) -> double
+
+    Compute the Poisson cdf with given lambda (mean) parameter:  P(N <= value; lambda).
+    The lambda parameter must be a positive real number (of type DOUBLE) and value must be a non-negative integer.
+
 .. function:: pow(x, p) -> double
 
     This is an alias for :func:`power`.

diff --git a/presto-main/src/main/java/com/facebook/presto/operator/scalar/MathFunctions.java b/presto-main/src/main/java/com/facebook/presto/operator/scalar/MathFunctions.java
@@ -33,6 +33,7 @@
 import org.apache.commons.math3.distribution.BetaDistribution;
 import org.apache.commons.math3.distribution.BinomialDistribution;
 import org.apache.commons.math3.distribution.ChiSquaredDistribution;
+import org.apache.commons.math3.distribution.PoissonDistribution;
 import org.apache.commons.math3.special.Erf;
 
 import java.math.BigDecimal;
@@ -761,6 +762,32 @@ public static double chiSquaredCdf(
         return distribution.cumulativeProbability(value);
     }
 
+    @Description("Inverse of Poisson cdf given lambda (mean) parameter and probability")
+    @ScalarFunction
+    @SqlType(StandardTypes.INTEGER)
+    public static long inversePoissonCdf(
+            @SqlType(StandardTypes.DOUBLE) double lambda,
+            @SqlType(StandardTypes.DOUBLE) double p)
+    {
+        checkCondition(p >= 0 && p < 1, INVALID_FUNCTION_ARGUMENT, "p must be in the interval [0, 1)");
+        checkCondition(lambda > 0, INVALID_FUNCTION_ARGUMENT, "lambda must be greater than 0");
+        PoissonDistribution distribution = new PoissonDistribution(lambda);
+        return distribution.inverseCumulativeProbability(p);
+    }
+
+    @Description("Poisson cdf given the lambda (mean) parameter and value")
+    @ScalarFunction
+    @SqlType(StandardTypes.DOUBLE)
+    public static double poissonCdf(
+            @SqlType(StandardTypes.DOUBLE) double lambda,
+            @SqlType(StandardTypes.INTEGER) long value)
+    {
+        checkCondition(value >= 0, INVALID_FUNCTION_ARGUMENT, "value must be a non-negative integer");
+        checkCondition(lambda > 0, INVALID_FUNCTION_ARGUMENT, "lambda must be greater than 0");
+        PoissonDistribution distribution = new PoissonDistribution(lambda);
+        return distribution.cumulativeProbability((int) value);
+    }
+
     @Description("round to nearest integer")
     @ScalarFunction("round")
     @SqlType(StandardTypes.TINYINT)

diff --git a/presto-main/src/test/java/com/facebook/presto/operator/scalar/TestMathFunctions.java b/presto-main/src/test/java/com/facebook/presto/operator/scalar/TestMathFunctions.java
@@ -1427,6 +1427,30 @@ public void testChiSquaredCdf()
         assertInvalidFunction("chi_squared_cdf(3, -10)", "value must non-negative");
     }
 
+    @Test
+    public void testInversePoissonCdf()
+    {
+        assertFunction("inverse_poisson_cdf(3, 0.0)", INTEGER, 0);
+        assertFunction("inverse_poisson_cdf(3, 0.3)", INTEGER, 2);
+        assertFunction("inverse_poisson_cdf(3, 0.95)", INTEGER, 6);
+        assertFunction("inverse_poisson_cdf(3, 0.99999999)", INTEGER, 17);
+
+        assertInvalidFunction("inverse_poisson_cdf(-3, 0.3)", "lambda must be greater than 0");
+        assertInvalidFunction("inverse_poisson_cdf(3, -0.1)", "p must be in the interval [0, 1)");
+        assertInvalidFunction("inverse_poisson_cdf(3, 1.1)", "p must be in the interval [0, 1)");
+        assertInvalidFunction("inverse_poisson_cdf(3, 1)", "p must be in the interval [0, 1)");
+    }
+
+    @Test
+    public void testPoissonCdf()
+    {
+        assertFunction("round(poisson_cdf(10, 0), 2)", DOUBLE, 0.0);
+        assertFunction("round(poisson_cdf(3, 5), 2)", DOUBLE, 0.92);
+
+        assertInvalidFunction("poisson_cdf(-3, 5)", "lambda must be greater than 0");
+        assertInvalidFunction("poisson_cdf(3, -10)", "value must be a non-negative integer");
+    }
+
     @Test
     public void testWilsonInterval()
     {