TDigest Value at Quantile Functions (facebookincubator#12529)

Summary:
Pull Request resolved: facebookincubator#12529

Add value_at_quantile , values_at_quantile functions

Reviewed By: Yuhta

Differential Revision: D70256593

fbshipit-source-id: 6860fe1f5d86caef5a9163dc09dd2c834279e9f5

Author: natashasehgal

velox/exec/fuzzer/PrestoQueryRunner.cpp

@@ -451,16 +451,17 @@ bool PrestoQueryRunner::isConstantExprSupported(
 bool PrestoQueryRunner::isSupported(const exec::FunctionSignature& signature) {
   // TODO: support queries with these types. Among the types below, hugeint is
   // not a native type in Presto, so fuzzer should not use it as the type of
-  // cast-to or constant literals. Hyperloglog can only be casted from varbinary
-  // and cannot be used as the type of constant literals. Interval year to month
-  // can only be casted from NULL and cannot be used as the type of constant
-  // literals. Json, Ipaddress, Ipprefix, and UUID require special handling,
-  // because Presto requires literals of these types to be valid, and doesn't
-  // allow creating HIVE columns of these types.
+  // cast-to or constant literals. Hyperloglog and TDigest can only be casted
+  // from varbinary and cannot be used as the type of constant literals.
+  // Interval year to month can only be casted from NULL and cannot be used as
+  // the type of constant literals. Json, Ipaddress, Ipprefix, and UUID require
+  // special handling, because Presto requires literals of these types to be
+  // valid, and doesn't allow creating HIVE columns of these types.
   return !(
       usesTypeName(signature, "interval year to month") ||
       usesTypeName(signature, "hugeint") ||
       usesTypeName(signature, "hyperloglog") ||
+      usesTypeName(signature, "tdigest") ||
       usesInputTypeName(signature, "ipaddress") ||
       usesInputTypeName(signature, "ipprefix") ||
       usesInputTypeName(signature, "uuid"));

velox/exec/tests/FunctionResolutionTest.cpp

@@ -24,6 +24,7 @@
 #include "velox/functions/prestosql/types/BingTileType.h"
 #include "velox/functions/prestosql/types/HyperLogLogType.h"
 #include "velox/functions/prestosql/types/JsonType.h"
+#include "velox/functions/prestosql/types/TDigestType.h"
 #include "velox/functions/prestosql/types/TimestampWithTimeZoneType.h"
 
 namespace {
@@ -289,6 +290,21 @@ TEST_F(FunctionResolutionTest, resolveCustomTypeHyperLogLog) {
   EXPECT_EQ(type->toString(), HYPERLOGLOG()->toString());
 }
 
+template <typename T>
+struct FuncTDigest {
+  VELOX_DEFINE_FUNCTION_TYPES(T);
+  bool call(out_type<SimpleTDigest<double>>&) {
+    return false;
+  }
+};
+
+TEST_F(FunctionResolutionTest, resolveCustomTypeTDigest) {
+  registerFunction<FuncTDigest, SimpleTDigest<double>>({"f_tdigest"});
+
+  auto type = exec::simpleFunctions().resolveFunction("f_tdigest", {})->type();
+  EXPECT_EQ(type->toString(), TDIGEST(DOUBLE())->toString());
+}
+
 template <typename T>
 struct FuncJson {
   VELOX_DEFINE_FUNCTION_TYPES(T);

velox/expression/fuzzer/ExpressionFuzzerTest.cpp

@@ -85,6 +85,12 @@ int main(int argc, char** argv) {
       "cardinality",
       "element_at",
       "width_bucket",
+      // Fuzzer and the underlying engine are confused about TDigest functions
+      // (since TDigest is a user defined type), and tries to pass a
+      // VARBINARY (since TDigest's implementation uses an
+      // alias to VARBINARY).
+      "value_at_quantile",
+      "values_at_quantiles",
       // Fuzzer cannot generate valid 'comparator' lambda.
       "array_sort(array(T),constant function(T,T,bigint)) -> array(T)",
       "split_to_map(varchar,varchar,varchar,function(varchar,varchar,varchar,varchar)) -> map(varchar,varchar)",

velox/functions/lib/TDigest.h

@@ -447,11 +447,15 @@ void TDigest<A>::mergeDeserialized(
     tdigest::detail::read(input, sum);
   }
   tdigest::detail::read(input, compression);
-  VELOX_CHECK_EQ(compression, compression_);
+  // If the TDigest is empty, set compression from TDigest being merged.
+  if (weights_.empty()) {
+    setCompression(compression);
+  }
   tdigest::detail::read(input, totalWeight);
   int32_t numNew;
   tdigest::detail::read(input, numNew);
   if (numNew > 0) {
+    VELOX_CHECK_EQ(compression, compression_);
     auto numOld = weights_.size();
     weights_.resize(numOld + numNew);
     auto* weights = weights_.data() + numOld;

velox/functions/prestosql/TDigestFunctions.h

@@ -0,0 +1,59 @@
+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include "velox/functions/Macros.h"
+#include "velox/functions/lib/TDigest.h"
+#include "velox/functions/prestosql/types/TDigestType.h"
+
+namespace facebook::velox::functions {
+
+template <typename T>
+struct ValueAtQuantileFunction {
+  VELOX_DEFINE_FUNCTION_TYPES(T);
+  FOLLY_ALWAYS_INLINE void call(
+      out_type<double>& result,
+      const arg_type<SimpleTDigest<double>>& input,
+      const arg_type<double>& quantile) {
+    TDigest<> digest;
+    std::vector<int16_t> positions;
+    digest.mergeDeserialized(positions, input.data());
+    digest.compress(positions);
+    result = digest.estimateQuantile(quantile);
+  }
+};
+
+template <typename T>
+struct ValuesAtQuantilesFunction {
+  VELOX_DEFINE_FUNCTION_TYPES(T);
+
+  FOLLY_ALWAYS_INLINE void call(
+      out_type<Array<double>>& result,
+      const arg_type<SimpleTDigest<double>>& input,
+      const arg_type<Array<double>>& quantiles) {
+    TDigest<> digest;
+    std::vector<int16_t> positions;
+    digest.mergeDeserialized(positions, input.data());
+    digest.compress(positions);
+    result.resize(quantiles.size());
+    for (size_t i = 0; i < quantiles.size(); ++i) {
+      result[i] = digest.estimateQuantile(quantiles[i].value());
+    }
+  }
+};
+
+} // namespace facebook::velox::functions

velox/functions/prestosql/registration/TDigestFunctionsRegistration.cpp

@@ -14,11 +14,22 @@
  * limitations under the License.
  */
 #include "velox/functions/Registerer.h"
+#include "velox/functions/prestosql/TDigestFunctions.h"
 #include "velox/functions/prestosql/types/TDigestRegistration.h"
-
+#include "velox/functions/prestosql/types/TDigestType.h"
 namespace facebook::velox::functions {
 
 void registerTDigestFunctions(const std::string& prefix) {
-  registerTDigestType();
+  facebook::velox::registerTDigestType();
+  registerFunction<
+      ValueAtQuantileFunction,
+      double,
+      SimpleTDigest<double>,
+      double>({prefix + "value_at_quantile"});
+  registerFunction<
+      ValuesAtQuantilesFunction,
+      Array<double>,
+      SimpleTDigest<double>,
+      Array<double>>({prefix + "values_at_quantiles"});
 }
 } // namespace facebook::velox::functions

velox/functions/prestosql/tests/TDigestFunctionsTest.cpp

@@ -0,0 +1,65 @@
+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include <folly/base64.h>
+#include "velox/functions/prestosql/tests/utils/FunctionBaseTest.h"
+#include "velox/functions/prestosql/types/TDigestRegistration.h"
+#include "velox/functions/prestosql/types/TDigestType.h"
+
+using namespace facebook::velox;
+using namespace facebook::velox::exec;
+using namespace facebook::velox::functions::test;
+
+class TDigestFunctionsTest : public FunctionBaseTest {
+ protected:
+  void SetUp() override {
+    FunctionBaseTest::SetUp();
+    registerTDigestType();
+  }
+
+ protected:
+  std::string decodeBase64(std::string_view input) {
+    std::string decoded(folly::base64DecodedSize(input), '\0');
+    folly::base64Decode(input, decoded.data());
+    return decoded;
+  }
+};
+
+TEST_F(TDigestFunctionsTest, valueAtQuantile) {
+  const TypePtr type = TDIGEST(DOUBLE());
+  const auto valueAtQuantile = [&](const std::optional<std::string>& input,
+                                   const std::optional<double>& quantile) {
+    return evaluateOnce<double>(
+        "value_at_quantile(c0, c1)", type, input, quantile);
+  };
+  const std::string input = decodeBase64(
+      "AQAAAAAAAADwPwAAAAAAABRAAAAAAAAALkAAAAAAAABZQAAAAAAAABRABQAAAAAAAAAAAPA/AAAAAAAA8D8AAAAAAADwPwAAAAAAAPA/AAAAAAAA8D8AAAAAAADwPwAAAAAAAABAAAAAAAAACEAAAAAAAAAQQAAAAAAAABRA");
+  ASSERT_EQ(1.0, valueAtQuantile(input, 0.1));
+  ASSERT_EQ(3.0, valueAtQuantile(input, 0.5));
+  ASSERT_EQ(5.0, valueAtQuantile(input, 0.9));
+  ASSERT_EQ(5.0, valueAtQuantile(input, 0.99));
+};
+
+TEST_F(TDigestFunctionsTest, valuesAtQuantiles) {
+  const TypePtr type = TDIGEST(DOUBLE());
+  const std::string input = decodeBase64(
+      "AQAAAAAAAADwPwAAAAAAABRAAAAAAAAALkAAAAAAAABZQAAAAAAAABRABQAAAAAAAAAAAPA/AAAAAAAA8D8AAAAAAADwPwAAAAAAAPA/AAAAAAAA8D8AAAAAAADwPwAAAAAAAABAAAAAAAAACEAAAAAAAAAQQAAAAAAAABRA");
+  VectorPtr arg0 = makeFlatVector<std::string>({input}, type);
+  VectorPtr arg1 = makeNullableArrayVector<double>({{0.1, 0.5, 0.9, 0.99}});
+  auto expected = makeNullableArrayVector<double>({{1.0, 3.0, 5.0, 5.0}});
+  auto result =
+      evaluate("values_at_quantiles(c0, c1)", makeRowVector({arg0, arg1}));
+  test::assertEqualVectors(expected, result);
+}