From 0adc62e3aaf647596ab476f6300a4165e60f6b91 Mon Sep 17 00:00:00 2001
From: Wills Feng <wills.feng@ibm.com>
Date: Mon, 15 Jul 2024 18:12:03 -0700
Subject: [PATCH] Add Presto inverse_laplace_cdf function (#9491)

Summary:
Resolves https://github.com/facebookincubator/velox/issues/5541
This is a followup for https://github.com/facebookincubator/velox/pull/5543

Pull Request resolved: https://github.com/facebookincubator/velox/pull/9491

Reviewed By: kagamiori, mbasmanova

Differential Revision: D59768749

Pulled By: pedroerp

fbshipit-source-id: 25b1e038570bf67e556df33db7ddcf7a1d4e0054
---
 velox/docs/functions/presto/math.rst          | 10 +++++--
 velox/functions/prestosql/Probability.h       | 17 +++++++++++
 ...lityTrigonometricFunctionsRegistration.cpp |  2 ++
 .../prestosql/tests/ProbabilityTest.cpp       | 29 +++++++++++++++++++
 4 files changed, 56 insertions(+), 2 deletions(-)

diff --git a/velox/docs/functions/presto/math.rst b/velox/docs/functions/presto/math.rst
index 300dd53802c0..c394182c6615 100644
--- a/velox/docs/functions/presto/math.rst
+++ b/velox/docs/functions/presto/math.rst
@@ -349,8 +349,14 @@ Probability Functions: inverse_cdf
 
 .. function:: inverse_cauchy_cdf(median, scale, p) -> double
 
-   Compute the inverse of the Cauchy cdf with given parameters ``median`` and ``scale`` (gamma) for the probability p.
-   The scale parameter must be a positive double. The probability ``p`` must be a double on the interval [0, 1].
+    Compute the inverse of the Cauchy cdf with given parameters ``median`` and ``scale`` (gamma) for the probability p.
+    The scale parameter must be a positive double. The probability ``p`` must be a double on the interval [0, 1].
+
+.. function:: inverse_laplace_cdf(mean, scale, p) -> double
+
+    Compute the inverse of the Laplace cdf with given ``mean`` and ``scale`` parameters for the cumulative probability (p): P(N < n).
+    The mean must be a real value and the scale must be a positive real value (both of type DOUBLE).
+    The probability ``p`` must lie on the interval [0, 1].
 
 ====================================
 Statistical Functions
diff --git a/velox/functions/prestosql/Probability.h b/velox/functions/prestosql/Probability.h
index 6b3fa48746eb..344fb44013f2 100644
--- a/velox/functions/prestosql/Probability.h
+++ b/velox/functions/prestosql/Probability.h
@@ -312,5 +312,22 @@ struct InverseCauchyCDFFunction {
   }
 };
 
+template <typename T>
+struct InverseLaplaceCDFFunction {
+  VELOX_DEFINE_FUNCTION_TYPES(T);
+  FOLLY_ALWAYS_INLINE void
+  call(double& result, double location, double scale, double p) {
+    VELOX_USER_CHECK_GT(scale, 0, "scale must be greater than 0");
+    VELOX_USER_CHECK(p >= 0 && p <= 1, "p must be in the interval [0, 1]")
+
+    if (std::isnan(location) || std::isinf(location)) {
+      result = std::numeric_limits<double>::quiet_NaN();
+    } else {
+      boost::math::laplace_distribution<> laplaceDist(location, scale);
+      result = boost::math::quantile(laplaceDist, p);
+    }
+  }
+};
+
 } // namespace
 } // namespace facebook::velox::functions
diff --git a/velox/functions/prestosql/registration/ProbabilityTrigonometricFunctionsRegistration.cpp b/velox/functions/prestosql/registration/ProbabilityTrigonometricFunctionsRegistration.cpp
index 8087120cccc8..c44489d3bcbf 100644
--- a/velox/functions/prestosql/registration/ProbabilityTrigonometricFunctionsRegistration.cpp
+++ b/velox/functions/prestosql/registration/ProbabilityTrigonometricFunctionsRegistration.cpp
@@ -79,6 +79,8 @@ void registerProbTrigFunctions(const std::string& prefix) {
       {prefix + "inverse_weibull_cdf"});
   registerFunction<InverseCauchyCDFFunction, double, double, double, double>(
       {prefix + "inverse_cauchy_cdf"});
+  registerFunction<InverseLaplaceCDFFunction, double, double, double, double>(
+      {prefix + "inverse_laplace_cdf"});
 }
 
 } // namespace
diff --git a/velox/functions/prestosql/tests/ProbabilityTest.cpp b/velox/functions/prestosql/tests/ProbabilityTest.cpp
index 361b469b6079..f85c9b9c140d 100644
--- a/velox/functions/prestosql/tests/ProbabilityTest.cpp
+++ b/velox/functions/prestosql/tests/ProbabilityTest.cpp
@@ -577,5 +577,34 @@ TEST_F(ProbabilityTest, inverseCauchyCDF) {
   EXPECT_EQ(invCauchyCDF(kDoubleMin, 2.0, 0.5), kDoubleMin);
 }
 
+TEST_F(ProbabilityTest, inverseLaplaceCDF) {
+  const auto inverseLaplaceCDF = [&](std::optional<double> location,
+                                     std::optional<double> scale,
+                                     std::optional<double> p) {
+    return evaluateOnce<double>(
+        "inverse_laplace_cdf(c0, c1, c2)", location, scale, p);
+  };
+
+  EXPECT_EQ(inverseLaplaceCDF(0.0, 1.0, 0.5), 0.0);
+  EXPECT_EQ(inverseLaplaceCDF(5.0, 2.0, 0.5), 5.0);
+
+  VELOX_ASSERT_THROW(
+      inverseLaplaceCDF(1.0, 1.0, kNan), "p must be in the interval [0, 1]");
+  VELOX_ASSERT_THROW(
+      inverseLaplaceCDF(1.0, 1.0, 2.0), "p must be in the interval [0, 1]");
+
+  EXPECT_EQ(inverseLaplaceCDF(10.0, kDoubleMax, 0.999999999999), kInf);
+  EXPECT_EQ(inverseLaplaceCDF(10.0, kDoubleMin, 0.000000000001), 10.0);
+  VELOX_ASSERT_THROW(
+      inverseLaplaceCDF(1.0, kNan, 0.5), "scale must be greater than 0");
+  VELOX_ASSERT_THROW(
+      inverseLaplaceCDF(1.0, -1.0, 0.5), "scale must be greater than 0");
+
+  EXPECT_THAT(inverseLaplaceCDF(kInf, 1.0, 0.5), IsNan());
+  EXPECT_THAT(inverseLaplaceCDF(kNan, 1.0, 0.5), IsNan());
+  EXPECT_THAT(inverseLaplaceCDF(kDoubleMax, 1.0, 0.5), kDoubleMax);
+  EXPECT_THAT(inverseLaplaceCDF(kDoubleMin, 1.0, 0.5), kDoubleMin);
+}
+
 } // namespace
 } // namespace facebook::velox