try

rui-mo · Mar 25, 2024 · 89e6cb8 · 89e6cb8
1 parent 9507d99
commit 89e6cb8
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 0 deletions.
diff --git a/velox/expression/tests/ExpressionFuzzer.cpp b/velox/expression/tests/ExpressionFuzzer.cpp
@@ -741,6 +741,9 @@ ExpressionFuzzer::ExpressionFuzzer(
   // Register function override (for cases where we want to restrict the types
   // or parameters we pass to functions).
   registerFuncOverride(&ExpressionFuzzer::generateSwitchArgs, "switch");
+  registerFuncOverride(
+      &ExpressionFuzzer::generateExtremeFunctionArgs, "greatest");
+  registerFuncOverride(&ExpressionFuzzer::generateExtremeFunctionArgs, "least");
 }
 
 void ExpressionFuzzer::getTicketsForFunctions() {
@@ -1017,6 +1020,39 @@ std::vector<core::TypedExprPtr> ExpressionFuzzer::generateSwitchArgs(
   return inputExpressions;
 }
 
+std::vector<core::TypedExprPtr> ExpressionFuzzer::generateExtremeFunctionArgs(
+    const CallableSignature& input) {
+  const auto argTypes = input.args;
+  VELOX_CHECK_EQ(
+      argTypes.size(),
+      1,
+      "Only one input is expected from the template signature.");
+  if (!argTypes[0]->isDecimal()) {
+    return generateArgs(input);
+  }
+
+  auto numVarArgs =
+      !input.variableArity ? 0 : rand32(0, options_.maxNumVarArgs);
+  std::vector<core::TypedExprPtr> inputExpressions;
+  inputExpressions.reserve(argTypes.size() + numVarArgs);
+  inputExpressions.emplace_back(
+      generateArg(argTypes.at(0), constantArgs.at(0)));
+
+  // Append varargs to the argument list.
+  for (int i = 0; i < numVarArgs; i++) {
+    size_t argClass = rand32(0, 1);
+    core::TypedExprPtr argExpr;
+    const auto argType = inputExpressions[0]->type();
+    if (argClass == kArgConstant) {
+      argExpr = generateArgConstant(argType);
+    } else {
+      argExpr = generateArgColumn(argType);
+    }
+    inputExpressions.emplace_back(argExpr);
+  }
+  return inputExpressions;
+}
+
 ExpressionFuzzer::FuzzedExpressionData ExpressionFuzzer::fuzzExpressions(
     const RowTypePtr& outType) {
   state.reset();

diff --git a/velox/expression/tests/ExpressionFuzzer.h b/velox/expression/tests/ExpressionFuzzer.h
@@ -271,6 +271,9 @@ class ExpressionFuzzer {
   std::vector<core::TypedExprPtr> generateSwitchArgs(
       const CallableSignature& input);
 
+  std::vector<core::TypedExprPtr> generateExtremeFunctionArgs(
+      const CallableSignature& input);
+
   TypePtr getConstrainedOutputType(
       const std::vector<core::TypedExprPtr>& args,
       const exec::FunctionSignature* signature);