TDigest Aggregate Fuzzer Test (#13301)

natashasehgal · facebook-github-bot · commit f4d10b3b0599 · 2025-06-17T11:12:37.000-07:00
Summary: Pull Request resolved: #13301 Differential Revision: D74505772
diff --git a/velox/functions/prestosql/fuzzer/AggregationFuzzerTest.cpp b/velox/functions/prestosql/fuzzer/AggregationFuzzerTest.cpp
@@ -37,6 +37,8 @@
 #include "velox/functions/prestosql/fuzzer/NoisyCountIfResultVerifier.h"
 #include "velox/functions/prestosql/fuzzer/QDigestAggInputGenerator.h"
 #include "velox/functions/prestosql/fuzzer/QDigestAggResultVerifier.h"
+#include "velox/functions/prestosql/fuzzer/TDigestAggregateInputGenerator.h"
+#include "velox/functions/prestosql/fuzzer/TDigestAggregateResultVerifier.h"
 #include "velox/functions/prestosql/registration/RegistrationFunctions.h"
 #include "velox/functions/prestosql/window/WindowFunctionsRegistration.h"
 #include "velox/vector/fuzzer/VectorFuzzer.h"
@@ -86,6 +88,7 @@ getCustomInputGenerators() {
       {"map_union_sum", std::make_shared<MapUnionSumInputGenerator>()},
       {"noisy_count_if_gaussian",
        std::make_shared<NoisyCountIfInputGenerator>()},
+      {"tdigest_agg", std::make_shared<TDigestAggregateInputGenerator>()},
   };
 }
 
@@ -154,6 +157,7 @@ int main(int argc, char** argv) {
   using facebook::velox::exec::test::NoisyCountIfResultVerifier;
   using facebook::velox::exec::test::QDigestAggResultVerifier;
   using facebook::velox::exec::test::setupReferenceQueryRunner;
+  using facebook::velox::exec::test::TDigestAggregateResultVerifier;
   using facebook::velox::exec::test::TransformResultVerifier;
 
   auto makeArrayVerifier = []() {
@@ -205,6 +209,7 @@ int main(int argc, char** argv) {
           {"sum_data_size_for_stats", nullptr},
           {"noisy_count_if_gaussian",
            std::make_shared<NoisyCountIfResultVerifier>()},
+          {"tdigest_agg", std::make_shared<TDigestAggregateResultVerifier>()},
       };
 
   using Runner = facebook::velox::exec::test::AggregationFuzzerRunner;
diff --git a/velox/functions/prestosql/fuzzer/TDigestAggregateInputGenerator.h b/velox/functions/prestosql/fuzzer/TDigestAggregateInputGenerator.h
@@ -0,0 +1,84 @@
+/*
+ * 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 <boost/random/uniform_int_distribution.hpp>
+#include <boost/random/uniform_real_distribution.hpp>
+#include <string>
+
+#include "velox/exec/fuzzer/InputGenerator.h"
+#include "velox/vector/FlatVector.h"
+#include "velox/vector/fuzzer/VectorFuzzer.h"
+
+namespace facebook::velox::exec::test {
+
+class TDigestAggregateInputGenerator : public InputGenerator {
+ public:
+  std::vector<VectorPtr> generate(
+      const std::vector<TypePtr>& types,
+      VectorFuzzer& fuzzer,
+      FuzzerGenerator& rng,
+      memory::MemoryPool* pool) override {
+    VELOX_CHECK_GE(types.size(), 1);
+    VELOX_CHECK_LE(types.size(), 3);
+
+    // Type checks
+    VELOX_CHECK(types[0]->isDouble());
+    if (types.size() > 1) {
+      VELOX_CHECK(types[1]->isBigint());
+    }
+    if (types.size() > 2) {
+      VELOX_CHECK(types[2]->isDouble());
+    }
+
+    std::vector<VectorPtr> inputs;
+    inputs.reserve(types.size());
+
+    // Values vector
+    auto valuesVector = fuzzer.fuzz(types[0]);
+    inputs.push_back(valuesVector);
+
+    // Weight is optional
+    if (types.size() > 1) {
+      auto weightsVector = fuzzer.fuzz(types[1]);
+      inputs.push_back(weightsVector);
+    }
+
+    // Compression is optional
+    if (types.size() > 2) {
+      const auto size = fuzzer.getOptions().vectorSize;
+      // Make sure to use the same value of 'compression' for all batches in a
+      // given Fuzzer iteration.
+      if (!compression_.has_value()) {
+        boost::random::uniform_real_distribution<double> dist(10.0, 1000.0);
+        compression_ = dist(rng);
+      }
+      inputs.push_back(BaseVector::createConstant(
+          DOUBLE(), compression_.value(), size, pool));
+    }
+
+    return inputs;
+  }
+
+  void reset() override {
+    compression_.reset();
+  }
+
+ private:
+  std::optional<double> compression_;
+};
+
+} // namespace facebook::velox::exec::test
diff --git a/velox/functions/prestosql/fuzzer/TDigestAggregateResultVerifier.h b/velox/functions/prestosql/fuzzer/TDigestAggregateResultVerifier.h
@@ -0,0 +1,194 @@
+/*
+ * 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/core/PlanNode.h"
+#include "velox/exec/fuzzer/ResultVerifier.h"
+#include "velox/exec/tests/utils/AssertQueryBuilder.h"
+#include "velox/exec/tests/utils/PlanBuilder.h"
+#include "velox/functions/lib/TDigest.h"
+#include "velox/vector/ComplexVector.h"
+
+namespace facebook::velox::exec::test {
+
+class TDigestAggregateResultVerifier : public ResultVerifier {
+ public:
+  bool supportsCompare() override {
+    return true;
+  }
+
+  bool supportsVerify() override {
+    return false;
+  }
+
+  void initialize(
+      const std::vector<RowVectorPtr>& /*input*/,
+      const std::vector<core::ExprPtr>& /*projections*/,
+      const std::vector<std::string>& groupingKeys,
+      const core::AggregationNode::Aggregate& aggregate,
+      const std::string& aggregateName) override {
+    keys_ = groupingKeys;
+    resultName_ = aggregateName;
+
+    // Check if the input type is double() which is the only valid one for
+    // TDigest
+    auto inputType = aggregate.call->inputs()[0]->type();
+    if (inputType->kind() != TypeKind::DOUBLE) {
+      VELOX_FAIL(
+          "TDigest only supports DOUBLE input type, got {}",
+          inputType->toString());
+    }
+  }
+
+  void initializeWindow(
+      const std::vector<RowVectorPtr>& /*input*/,
+      const std::vector<core::ExprPtr>& /*projections*/,
+      const std::vector<std::string>& /*partitionByKeys*/,
+      const std::vector<SortingKeyAndOrder>& /*sortingKeysAndOrders*/,
+      const core::WindowNode::Function& function,
+      const std::string& /*frame*/,
+      const std::string& windowName) override {
+    keys_ = {"row_number"};
+    resultName_ = windowName;
+
+    // Check input type is double
+    auto inputType = function.functionCall->inputs()[0]->type();
+    if (inputType->kind() != TypeKind::DOUBLE) {
+      VELOX_FAIL(
+          "TDigest only supports DOUBLE input type, got {}",
+          inputType->toString());
+    }
+
+    // Check return type is varbinary
+    auto returnType = function.functionCall->type();
+    if (returnType->kind() != TypeKind::VARBINARY) {
+      VELOX_FAIL(
+          "TDigest return type must be VARBINARY, got {}",
+          returnType->toString());
+    }
+  }
+
+  bool compare(const RowVectorPtr& result, const RowVectorPtr& altResult)
+      override {
+    VELOX_CHECK_EQ(result->size(), altResult->size());
+
+    auto projection = keys_;
+    projection.push_back(resultName_);
+
+    auto planNodeIdGenerator = std::make_shared<core::PlanNodeIdGenerator>();
+    auto builder = PlanBuilder(planNodeIdGenerator).values({result});
+    if (!keys_.empty()) {
+      builder = builder.orderBy(keys_, false);
+    }
+    auto sortByKeys = builder.project(projection).planNode();
+    auto sortedResult =
+        AssertQueryBuilder(sortByKeys).copyResults(result->pool());
+
+    builder = PlanBuilder(planNodeIdGenerator).values({altResult});
+    if (!keys_.empty()) {
+      builder = builder.orderBy(keys_, false);
+    }
+    sortByKeys = builder.project(projection).planNode();
+    auto sortedAltResult =
+        AssertQueryBuilder(sortByKeys).copyResults(altResult->pool());
+
+    VELOX_CHECK_EQ(sortedResult->size(), sortedAltResult->size());
+    auto size = sortedResult->size();
+    for (auto i = 0; i < size; i++) {
+      auto resultIsNull = sortedResult->childAt(resultName_)->isNullAt(i);
+      auto altResultIsNull = sortedAltResult->childAt(resultName_)->isNullAt(i);
+      if (resultIsNull || altResultIsNull) {
+        VELOX_CHECK(resultIsNull && altResultIsNull);
+        continue;
+      }
+
+      auto resultValue = sortedResult->childAt(resultName_)
+                             ->as<SimpleVector<StringView>>()
+                             ->valueAt(i);
+      auto altResultValue = sortedAltResult->childAt(resultName_)
+                                ->as<SimpleVector<StringView>>()
+                                ->valueAt(i);
+      if (resultValue == altResultValue) {
+        continue;
+      } else {
+        checkEquivalentTDigest(resultValue, altResultValue);
+      }
+    }
+    return true;
+  }
+
+  bool verify(const RowVectorPtr& /*result*/) override {
+    VELOX_UNSUPPORTED();
+  }
+
+  void reset() override {
+    keys_.clear();
+    resultName_.clear();
+  }
+
+ private:
+  void checkEquivalentTDigest(
+      const StringView& result,
+      const StringView& altResult) {
+    // Create TDigests from serialized data
+    facebook::velox::functions::TDigest<> resultTdigest;
+    facebook::velox::functions::TDigest<> altResultTdigest;
+    std::vector<int16_t> positions;
+
+    try {
+      resultTdigest.mergeDeserialized(positions, result.data());
+      resultTdigest.compress(positions);
+
+      positions.clear();
+      altResultTdigest.mergeDeserialized(positions, altResult.data());
+      altResultTdigest.compress(positions);
+    } catch (const std::exception& e) {
+      VELOX_FAIL("Failed to deserialize TDigest: {}", e.what());
+    }
+
+    // Compare TDigest values at specific quantiles
+    for (auto quantile : kQuantiles) {
+      double resultQuantile = resultTdigest.estimateQuantile(quantile);
+      double altResultQuantile = altResultTdigest.estimateQuantile(quantile);
+
+      VELOX_CHECK(
+          std::abs(resultQuantile - altResultQuantile) <= kError,
+          "TDigest quantile values differ at {}: {} vs {}",
+          quantile,
+          resultQuantile,
+          altResultQuantile);
+    }
+  }
+
+  static constexpr double kError = 0.0001;
+
+  static constexpr double kQuantiles[] = {
+      0.01,
+      0.05,
+      0.1,
+      0.25,
+      0.50,
+      0.75,
+      0.9,
+      0.95,
+      0.99,
+  };
+
+  std::vector<std::string> keys_;
+  std::string resultName_;
+};
+
+} // namespace facebook::velox::exec::test
diff --git a/velox/functions/prestosql/fuzzer/WindowFuzzerTest.cpp b/velox/functions/prestosql/fuzzer/WindowFuzzerTest.cpp
@@ -29,6 +29,8 @@
 #include "velox/functions/prestosql/fuzzer/MinMaxInputGenerator.h"
 #include "velox/functions/prestosql/fuzzer/QDigestAggInputGenerator.h"
 #include "velox/functions/prestosql/fuzzer/QDigestAggResultVerifier.h"
+#include "velox/functions/prestosql/fuzzer/TDigestAggregateInputGenerator.h"
+#include "velox/functions/prestosql/fuzzer/TDigestAggregateResultVerifier.h"
 #include "velox/functions/prestosql/fuzzer/WindowOffsetInputGenerator.h"
 #include "velox/functions/prestosql/registration/RegistrationFunctions.h"
 #include "velox/functions/prestosql/window/WindowFunctionsRegistration.h"
@@ -76,6 +78,7 @@ getCustomInputGenerators() {
       {"approx_set", std::make_shared<ApproxDistinctInputGenerator>()},
       {"approx_percentile", std::make_shared<ApproxPercentileInputGenerator>()},
       {"qdigest_agg", std::make_shared<QDigestAggInputGenerator>()},
+      {"tdigest_agg", std::make_shared<TDigestAggregateInputGenerator>()},
       {"lead", std::make_shared<WindowOffsetInputGenerator>(1)},
       {"lag", std::make_shared<WindowOffsetInputGenerator>(1)},
       {"nth_value", std::make_shared<WindowOffsetInputGenerator>(1)},
@@ -143,6 +146,7 @@ int main(int argc, char** argv) {
   using facebook::velox::exec::test::ApproxPercentileResultVerifier;
   using facebook::velox::exec::test::AverageResultVerifier;
   using facebook::velox::exec::test::QDigestAggResultVerifier;
+  using facebook::velox::exec::test::TDigestAggregateResultVerifier;
 
   static const std::unordered_map<
       std::string,
@@ -155,6 +159,7 @@ int main(int argc, char** argv) {
            std::make_shared<ApproxPercentileResultVerifier>()},
           {"approx_most_frequent", nullptr},
           {"qdigest_agg", std::make_shared<QDigestAggResultVerifier>()},
+          {"tdigest_agg", std::make_shared<TDigestAggregateResultVerifier>()},
           {"merge", nullptr},
           // Semantically inconsistent functions
           {"skewness", nullptr},
@@ -192,6 +197,7 @@ int main(int argc, char** argv) {
       "min_by",
       "multimap_agg",
       "qdigest_agg",
+      "tdigest_agg",
   };
 
   using Runner = facebook::velox::exec::test::WindowFuzzerRunner;
