feat: Add inverse_binomial_cdf (facebookincubator#12983)

Summary:
Pull Request resolved: facebookincubator#12983

feat: Add inverse_binomial_cdf

Reviewed By: amitkdutta

Differential Revision: D72749280

fbshipit-source-id: ef11432ee80f2034a45411512c4d025ef55a90a0

Author: skyelves

velox/expression/fuzzer/ExpressionFuzzerTest.cpp

@@ -221,6 +221,7 @@ int main(int argc, char** argv) {
         "map_keys_by_top_n_values", // requires
                                     // https://github.com/prestodb/presto/pull/24570
         "inverse_gamma_cdf", // https://github.com/facebookincubator/velox/issues/12918
+        "inverse_binomial_cdf", // https://github.com/facebookincubator/velox/issues/12981
     });
 
     referenceQueryRunner = std::make_shared<PrestoQueryRunner>(

velox/functions/prestosql/Probability.h

@@ -355,5 +355,33 @@ struct InverseGammaCDFFunction {
   }
 };
 
+template <typename T>
+struct InverseBinomialCDFFunction {
+  VELOX_DEFINE_FUNCTION_TYPES(T);
+
+  FOLLY_ALWAYS_INLINE void call(
+      int32_t& result,
+      int32_t numberOfTrials,
+      double successProbability,
+      double p) {
+    static constexpr double kInf = std::numeric_limits<double>::infinity();
+
+    VELOX_USER_CHECK(
+        (p >= 0) && (p <= 1) && (p != kInf),
+        "inverseBinomialCdf Function: p must be in the interval [0, 1]");
+    VELOX_USER_CHECK(
+        (successProbability >= 0) && (successProbability <= 1) &&
+            (successProbability != kInf),
+        "inverseBinomialCdf Function: successProbability must be in the interval [0, 1]");
+    VELOX_USER_CHECK(
+        numberOfTrials > 0,
+        "inverseBinomialCdf Function: numberOfTrials must be greater than 0");
+
+    boost::math::binomial_distribution<> dist(
+        numberOfTrials, successProbability);
+    result = static_cast<int32_t>(boost::math::quantile(dist, p));
+  }
+};
+
 } // namespace
 } // namespace facebook::velox::functions

velox/functions/prestosql/registration/ProbabilityTrigonometricFunctionsRegistration.cpp

@@ -81,6 +81,12 @@ void registerProbTrigFunctions(const std::string& prefix) {
       {prefix + "inverse_laplace_cdf"});
   registerFunction<InverseGammaCDFFunction, double, double, double, double>(
       {prefix + "inverse_gamma_cdf"});
+  registerFunction<
+      InverseBinomialCDFFunction,
+      int32_t,
+      int32_t,
+      double,
+      double>({prefix + "inverse_binomial_cdf"});
 }
 
 } // namespace

velox/functions/prestosql/tests/ProbabilityTest.cpp

@@ -674,5 +674,56 @@ TEST_F(ProbabilityTest, invGammaCDF) {
       "inverseGammaCdf Function: p must be in the interval [0, 1]");
 }
 
+TEST_F(ProbabilityTest, invBinomialCDF) {
+  const auto invBinomialCDF = [&](std::optional<int32_t> numberOfTrials,
+                                  std::optional<double> successProbability,
+                                  std::optional<double> p) {
+    return evaluateOnce<int32_t>(
+        "inverse_binomial_cdf(c0, c1, c2)",
+        numberOfTrials,
+        successProbability,
+        p);
+  };
+
+  EXPECT_EQ(0, invBinomialCDF(20, 0.5, 0.0));
+  EXPECT_EQ(10, invBinomialCDF(20, 0.5, 0.5));
+  EXPECT_EQ(20, invBinomialCDF(20, 0.5, 1.0));
+  EXPECT_EQ(INT32_MAX, invBinomialCDF(INT32_MAX, 0.5, 1));
+  EXPECT_EQ(611204, invBinomialCDF(1223340, 0.5, 0.2));
+
+  EXPECT_EQ(std::nullopt, invBinomialCDF(std::nullopt, 1, 1));
+  EXPECT_EQ(std::nullopt, invBinomialCDF(1, std::nullopt, 1));
+  EXPECT_EQ(std::nullopt, invBinomialCDF(1, 1, std::nullopt));
+
+  VELOX_ASSERT_THROW(
+      invBinomialCDF(5, -0.5, 0.3),
+      "inverseBinomialCdf Function: successProbability must be in the interval [0, 1]");
+  VELOX_ASSERT_THROW(
+      invBinomialCDF(5, 1.5, 0.3),
+      "inverseBinomialCdf Function: successProbability must be in the interval [0, 1]");
+  VELOX_ASSERT_THROW(
+      invBinomialCDF(5, 0.5, -3.0),
+      "inverseBinomialCdf Function: p must be in the interval [0, 1]");
+  VELOX_ASSERT_THROW(
+      invBinomialCDF(5, 0.5, 3.0),
+      "inverseBinomialCdf Function: p must be in the interval [0, 1]");
+  VELOX_ASSERT_THROW(
+      invBinomialCDF(-5, 0.5, 0.3),
+      "inverseBinomialCdf Function: numberOfTrials must be greater than 0");
+
+  VELOX_ASSERT_THROW(
+      invBinomialCDF(1, kInf, 0.5),
+      "inverseBinomialCdf Function: successProbability must be in the interval [0, 1]");
+  VELOX_ASSERT_THROW(
+      invBinomialCDF(1, kNan, 0.5),
+      "inverseBinomialCdf Function: successProbability must be in the interval [0, 1]");
+  VELOX_ASSERT_THROW(
+      invBinomialCDF(1, 0.5, kInf),
+      "inverseBinomialCdf Function: p must be in the interval [0, 1]");
+  VELOX_ASSERT_THROW(
+      invBinomialCDF(1, 0.5, kNan),
+      "inverseBinomialCdf Function: p must be in the interval [0, 1]");
+}
+
 } // namespace
 } // namespace facebook::velox