Add distinct aggregations to AggregationFuzzer

velox/exec/DistinctAggregations.cpp

@@ -237,6 +237,8 @@ std::unique_ptr<DistinctAggregations> DistinctAggregations::create(
     case TypeKind::TIMESTAMP:
       return std::make_unique<TypedDistinctAggregations<Timestamp>>(
           aggregates, inputType, pool);
+    case TypeKind::VARBINARY:
+      [[fallthrough]];
     case TypeKind::VARCHAR:
       return std::make_unique<TypedDistinctAggregations<StringView>>(
           aggregates, inputType, pool);

velox/exec/fuzzer/AggregationFuzzer.cpp

@@ -23,6 +23,7 @@
 #include "velox/exec/tests/utils/PlanBuilder.h"
 #include "velox/exec/tests/utils/TempDirectoryPath.h"
 
+#include <algorithm>
 #include "velox/exec/PartitionFunction.h"
 #include "velox/exec/fuzzer/AggregationFuzzerBase.h"
 #include "velox/expression/tests/utils/FuzzerToolkit.h"
@@ -136,6 +137,14 @@ class AggregationFuzzer : public AggregationFuzzerBase {
       const std::vector<std::string>& masks,
       const std::vector<RowVectorPtr>& input);
 
+  bool verifyDistinctAggregation(
+      const std::vector<std::string>& groupingKeys,
+      const std::vector<std::string>& aggregates,
+      const std::vector<std::string>& masks,
+      const std::vector<RowVectorPtr>& input,
+      bool customVerification,
+      const std::vector<std::shared_ptr<ResultVerifier>>& customVerifiers);
+
   void verifyAggregation(const std::vector<PlanWithSplits>& plans);
 
   static bool hasPartialGroupBy(const core::PlanNodePtr& plan) {
@@ -319,11 +328,39 @@ void AggregationFuzzer::go() {
 
       std::vector<TypePtr> types;
       std::vector<std::string> names;
-
       auto groupingKeys = generateKeys("g", names, types);
+
+      std::vector<TypePtr> payloadTypes;
+      std::vector<std::string> payloadNames;
+      auto payloads = generateKeys("c", payloadNames, payloadTypes);
+
+      names.insert(
+          names.end(),
+          std::make_move_iterator(payloadNames.begin()),
+          std::make_move_iterator(payloadNames.end()));
+      types.insert(
+          types.end(),
+          std::make_move_iterator(payloadTypes.begin()),
+          std::make_move_iterator(payloadTypes.end()));
+
+      payloadNames.clear();
+      payloadTypes.clear();
+
       auto input = generateInputData(names, types, std::nullopt);
 
       verifyAggregation(groupingKeys, {}, {}, input, false, {});
+      std::shared_ptr<ResultVerifier> customVerifier;
+      customVerifier = customVerificationFunctions_.at("array_agg");
+      verifyDistinctAggregation(
+          groupingKeys,
+          {"array_agg(distinct c0)"},
+          {},
+          input,
+          true,
+          {customVerifier});
+      verifyDistinctAggregation(
+          groupingKeys, {"count(distinct c0)"}, {}, input, false, {});
+
     } else {
       // Pick a random signature.
       auto signatureWithStats = pickSignature();
@@ -915,6 +952,111 @@ bool AggregationFuzzer::verifySortedAggregation(
   return resultOrError.exceptionPtr != nullptr;
 }
 
+bool AggregationFuzzer::verifyDistinctAggregation(
+    const std::vector<std::string>& groupingKeys,
+    const std::vector<std::string>& aggregates,
+    const std::vector<std::string>& masks,
+    const std::vector<RowVectorPtr>& input,
+    bool customVerification,
+    const std::vector<std::shared_ptr<ResultVerifier>>& customVerifiers) {
+  auto firstPlan = PlanBuilder()
+                       .values(input)
+                       .singleAggregation(groupingKeys, aggregates, masks)
+                       .planNode();
+
+  if (customVerification) {
+    const auto& aggregationNode =
+        std::dynamic_pointer_cast<const core::AggregationNode>(firstPlan);
+
+    for (auto i = 0; i < customVerifiers.size(); ++i) {
+      auto& verifier = customVerifiers[i];
+      if (verifier == nullptr) {
+        continue;
+      }
+
+      verifier->initialize(
+          input,
+          groupingKeys,
+          aggregationNode->aggregates()[i],
+          aggregationNode->aggregateNames()[i]);
+    }
+  }
+
+  SCOPE_EXIT {
+    if (customVerification) {
+      resetCustomVerifiers(customVerifiers);
+    }
+  };
+
+  // Create all the plans upfront.
+  std::vector<PlanWithSplits> plans;
+  auto directory = exec::test::TempDirectoryPath::create();
+  const auto inputRowType = asRowType(input[0]->type());
+
+  if (isTableScanSupported(inputRowType)) {
+    directory = exec::test::TempDirectoryPath::create();
+    auto splits = makeSplits(input, directory->path);
+    plans.push_back(
+        {PlanBuilder()
+             .tableScan(inputRowType)
+             .singleAggregation(groupingKeys, aggregates, masks)
+             .planNode(),
+         splits});
+  }
+
+  if (persistAndRunOnce_) {
+    persistReproInfo(plans, reproPersistPath_);
+  }
+
+  try {
+    auto resultOrError = execute(firstPlan);
+    if (resultOrError.exceptionPtr) {
+      ++stats_.numFailed;
+    }
+
+    // TODO Use ResultVerifier::compare API to compare Velox results with
+    // reference DB results once reference query runner is updated to return
+    // results as Velox vectors.
+    std::optional<MaterializedRowMultiset> expectedResult;
+    if (resultOrError.result != nullptr) {
+      if (!customVerification) {
+        auto referenceResult = computeReferenceResults(firstPlan, input);
+        updateReferenceQueryStats(referenceResult.second);
+        expectedResult = referenceResult.first;
+      } else {
+        ++stats_.numVerificationSkipped;
+
+        for (auto& verifier : customVerifiers) {
+          if (verifier != nullptr && verifier->supportsVerify()) {
+            VELOX_CHECK(
+                verifier->verify(resultOrError.result),
+                "Aggregation results failed custom verification");
+          }
+        }
+      }
+    }
+
+    if (expectedResult && resultOrError.result) {
+      ++stats_.numVerified;
+      VELOX_CHECK(
+          assertEqualResults(
+              expectedResult.value(),
+              firstPlan->outputType(),
+              {resultOrError.result}),
+          "Velox and reference DB results don't match");
+      LOG(INFO) << "Verified results against reference DB";
+    }
+
+    testPlans(plans, customVerification, customVerifiers, resultOrError, 1);
+
+    return resultOrError.exceptionPtr != nullptr;
+  } catch (...) {
+    if (!reproPersistPath_.empty()) {
+      persistReproInfo(plans, reproPersistPath_);
+    }
+    throw;
+  }
+}
 // verifyAggregation(std::vector<core::PlanNodePtr> plans) is tied to plan
 // created by previous verifyAggregation function. Changes in nodes there will
 // require corresponding changes here.