Velox: tdigest_agg (#12315)

Summary: Add aggregate function tdigest_agg() in Velox Differential Revision: D69331895
facebookincubator · Feb 14, 2025 · c87b683 · c87b683
1 parent 2d079bc
commit c87b683
Show file tree

Hide file tree

Showing 8 changed files with 445 additions and 2 deletions.
diff --git a/velox/functions/lib/CMakeLists.txt b/velox/functions/lib/CMakeLists.txt
@@ -39,6 +39,7 @@ velox_add_library(
   StringEncodingUtils.cpp
   SubscriptUtil.cpp
   TimeUtils.cpp
+  TDigest.h
   Utf8Utils.cpp)
 
 velox_link_libraries(

diff --git a/velox/functions/lib/TDigest.h b/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;

diff --git a/velox/functions/prestosql/aggregates/AggregateNames.h b/velox/functions/prestosql/aggregates/AggregateNames.h
@@ -73,6 +73,7 @@ const char* const kStdDev = "stddev"; // Alias for stddev_samp.
 const char* const kStdDevPop = "stddev_pop";
 const char* const kStdDevSamp = "stddev_samp";
 const char* const kSum = "sum";
+const char* const kTDigestAgg = "tdigest_agg";
 const char* const kVariance = "variance"; // Alias for var_samp.
 const char* const kVarPop = "var_pop";
 const char* const kVarSamp = "var_samp";

diff --git a/velox/functions/prestosql/aggregates/CMakeLists.txt b/velox/functions/prestosql/aggregates/CMakeLists.txt
@@ -46,6 +46,7 @@ velox_add_library(
   SetAggregates.cpp
   SumAggregate.cpp
   SumDataSizeForStatsAggregate.cpp
+  TDigestAggregate.cpp
   VarianceAggregates.cpp)
 
 velox_link_libraries(

diff --git a/velox/functions/prestosql/aggregates/RegisterAggregateFunctions.cpp b/velox/functions/prestosql/aggregates/RegisterAggregateFunctions.cpp
@@ -16,6 +16,7 @@
 #include "velox/functions/prestosql/aggregates/RegisterAggregateFunctions.h"
 #include "velox/exec/Aggregate.h"
 #include "velox/functions/prestosql/types/JsonType.h"
+#include "velox/functions/prestosql/types/TDigestType.h"
 
 namespace facebook::velox::aggregate::prestosql {
 
@@ -150,13 +151,18 @@ extern void registerVarianceAggregates(
     const std::string& prefix,
     bool withCompanionFunctions,
     bool overwrite);
+extern void registerTDigestAggregate(
+    const std::string& prefix,
+    bool withCompanionFunctions,
+    bool overwrite);
 
 void registerAllAggregateFunctions(
     const std::string& prefix,
     bool withCompanionFunctions,
     bool onlyPrestoSignatures,
     bool overwrite) {
   registerJsonType();
+  registerTDigestType();
   registerApproxDistinctAggregates(prefix, withCompanionFunctions, overwrite);
   registerApproxMostFrequentAggregate(
       prefix, withCompanionFunctions, overwrite);
@@ -194,6 +200,7 @@ void registerAllAggregateFunctions(
   registerSetUnionAggregate(prefix, withCompanionFunctions, overwrite);
   registerSumAggregate(prefix, withCompanionFunctions, overwrite);
   registerVarianceAggregates(prefix, withCompanionFunctions, overwrite);
+  registerTDigestAggregate(prefix, withCompanionFunctions, overwrite);
 }
 
 extern void registerCountDistinctAggregate(