support decimal in expression fuzzer test

rui-mo · Mar 26, 2024 · 4f17e1e · 4f17e1e
1 parent 3aa020d
commit 4f17e1e
Show file tree

Hide file tree

Showing 13 changed files with 424 additions and 39 deletions.
diff --git a/velox/expression/ConstantExpr.cpp b/velox/expression/ConstantExpr.cpp
@@ -162,6 +162,7 @@ void appendSqlLiteral(
     case TypeKind::TINYINT:
     case TypeKind::SMALLINT:
     case TypeKind::BIGINT:
+    case TypeKind::HUGEINT:
     case TypeKind::TIMESTAMP:
     case TypeKind::REAL:
     case TypeKind::DOUBLE:

diff --git a/velox/expression/ReverseSignatureBinder.cpp b/velox/expression/ReverseSignatureBinder.cpp
@@ -36,9 +36,6 @@ bool ReverseSignatureBinder::hasConstrainedIntegerVariable(
 }
 
 bool ReverseSignatureBinder::tryBind() {
-  if (hasConstrainedIntegerVariable(signature_.returnType())) {
-    return false;
-  }
   return SignatureBinderBase::tryBind(signature_.returnType(), returnType_);
 }
 

diff --git a/velox/expression/tests/ArgumentTypeFuzzerTest.cpp b/velox/expression/tests/ArgumentTypeFuzzerTest.cpp
@@ -59,6 +59,42 @@ class ArgumentTypeFuzzerTest : public testing::Test {
     }
   }
 
+  void testFuzzingDecimalSuccess(
+      const std::shared_ptr<exec::FunctionSignature>& signature,
+      int32_t expectedArguments,
+      std::optional<TypeKind> outputKind = std::nullopt) {
+    std::mt19937 seed{0};
+    ArgumentTypeFuzzer fuzzer{*signature, seed};
+    ASSERT_TRUE(fuzzer.fuzzArgumentTypes(kMaxVariadicArgs));
+
+    auto& argumentTypes = fuzzer.argumentTypes();
+    ASSERT_LE(argumentTypes.size(), expectedArguments);
+
+    auto& argumentSignatures = signature->argumentTypes();
+    int i;
+    for (i = 0; i < expectedArguments; ++i) {
+      ASSERT_TRUE(argumentTypes[i]->isDecimal())
+          << "at " << i
+          << ": Expected decimal. Got: " << argumentTypes[i]->toString();
+    }
+
+    if (i < argumentTypes.size()) {
+      ASSERT_TRUE(signature->variableArity());
+      ASSERT_LE(
+          argumentTypes.size() - argumentSignatures.size(), kMaxVariadicArgs);
+      for (int j = i; j < argumentTypes.size(); ++j) {
+        ASSERT_TRUE(argumentTypes[j]->equivalent(*argumentTypes[i - 1]));
+      }
+    }
+
+    const auto outputType = fuzzer.fuzzReturnType();
+    if (outputKind.has_value()) {
+      ASSERT_TRUE(outputType->kind() == outputKind);
+    } else {
+      ASSERT_TRUE(outputType->isDecimal());
+    }
+  }
+
   void testFuzzingFailure(
       const std::shared_ptr<exec::FunctionSignature>& signature,
       const TypePtr& returnType) {
@@ -222,9 +258,36 @@ TEST_F(ArgumentTypeFuzzerTest, any) {
   ASSERT_TRUE(argumentTypes[0] != nullptr);
 }
 
-TEST_F(ArgumentTypeFuzzerTest, unsupported) {
-  // Constraints on the return type is not supported.
-  auto signature =
+TEST_F(ArgumentTypeFuzzerTest, decimal) {
+  auto signature = exec::FunctionSignatureBuilder()
+                       .integerVariable("a_scale")
+                       .integerVariable("a_precision")
+                       .returnType("boolean")
+                       .argumentType("decimal(a_precision, a_scale)")
+                       .argumentType("decimal(a_precision, a_scale)")
+                       .argumentType("decimal(a_precision, a_scale)")
+                       .build();
+
+  testFuzzingDecimalSuccess(signature, 3, TypeKind::BOOLEAN);
+
+  signature =
+      exec::FunctionSignatureBuilder()
+          .integerVariable("a_precision")
+          .integerVariable("a_scale")
+          .integerVariable("b_precision")
+          .integerVariable("b_scale")
+          .integerVariable(
+              "r_precision",
+              "min(38, max(a_precision - a_scale, b_precision - b_scale) + max(a_scale, b_scale) + 1)")
+          .integerVariable("r_scale", "max(a_scale, b_scale)")
+          .returnType("DECIMAL(r_precision, r_scale)")
+          .argumentType("DECIMAL(a_precision, a_scale)")
+          .argumentType("DECIMAL(b_precision, b_scale)")
+          .build();
+
+  testFuzzingDecimalSuccess(signature, 2);
+
+  signature =
       exec::FunctionSignatureBuilder()
           .integerVariable("a_scale")
           .integerVariable("b_scale")
@@ -239,7 +302,31 @@ TEST_F(ArgumentTypeFuzzerTest, unsupported) {
           .argumentType("decimal(b_precision, b_scale)")
           .build();
 
-  testFuzzingFailure(signature, DECIMAL(13, 6));
+  testFuzzingDecimalSuccess(signature, 2, TypeKind::ROW);
+
+  signature = exec::FunctionSignatureBuilder()
+                  .integerVariable("i1")
+                  .integerVariable("i2")
+                  .integerVariable("i5")
+                  .integerVariable("i6")
+                  .integerVariable(
+                      "i3", "min(38, max(i1 - i5, i2 - i6) + max(i5, i6) + 1)")
+                  .integerVariable("i7", "max(i5, i6)")
+                  .returnType("decimal(i3,i7)")
+                  .argumentType("decimal(i1,i5)")
+                  .argumentType("decimal(i2,i6)")
+                  .build();
+  testFuzzingDecimalSuccess(signature, 2);
+
+  signature = exec::FunctionSignatureBuilder()
+                  .integerVariable("i1")
+                  .integerVariable("i5")
+                  .returnType("boolean")
+                  .argumentType("decimal(i1,i5)")
+                  .argumentType("decimal(i1,i5)")
+                  .argumentType("decimal(i1,i5)")
+                  .build();
+  testFuzzingDecimalSuccess(signature, 3, TypeKind::BOOLEAN);
 }
 
 TEST_F(ArgumentTypeFuzzerTest, lambda) {