feat: Support variadic function in CometFuzz (#2682)

* feat: Support variadic function in CometFuzz

* Check input type array has single element

Author: manuzhang

fuzz-testing/src/main/scala/org/apache/comet/fuzz/Meta.scala

@@ -44,7 +44,7 @@ case class SparkMapType(keyType: SparkType, valueType: SparkType) extends SparkT
 case class SparkStructType(fields: Seq[SparkType]) extends SparkType
 case object SparkAnyType extends SparkType
 
-case class FunctionSignature(inputTypes: Seq[SparkType])
+case class FunctionSignature(inputTypes: Seq[SparkType], varArgs: Boolean = false)
 
 case class Function(name: String, signatures: Seq[FunctionSignature])
 
@@ -65,11 +65,20 @@ object Meta {
     (DataTypes.StringType, 0.2),
     (DataTypes.BinaryType, 0.1))
 
-  private def createFunctionWithInputTypes(name: String, inputs: Seq[SparkType]): Function = {
-    Function(name, Seq(FunctionSignature(inputs)))
+  private def createFunctionWithInputTypes(
+      name: String,
+      inputs: Seq[SparkType],
+      varArgs: Boolean = false): Function = {
+    Function(name, Seq(FunctionSignature(inputs, varArgs)))
+    createFunctions(name, Seq(FunctionSignature(inputs, varArgs)))
   }
 
   private def createFunctions(name: String, signatures: Seq[FunctionSignature]): Function = {
+    signatures.foreach { s =>
+      assert(
+        !s.varArgs || s.inputTypes.length == 1,
+        s"Variadic function $s must have exactly one input type")
+    }
     Function(name, signatures)
   }
 
@@ -126,13 +135,11 @@ object Meta {
         SparkTypeOneOf(
           Seq(
             SparkStringType,
-            SparkArrayType(
-              SparkTypeOneOf(Seq(SparkStringType, SparkNumericType, SparkBinaryType))))),
-        SparkTypeOneOf(
-          Seq(
-            SparkStringType,
-            SparkArrayType(
-              SparkTypeOneOf(Seq(SparkStringType, SparkNumericType, SparkBinaryType))))))),
+            SparkBinaryType,
+            SparkArrayType(SparkStringType),
+            SparkArrayType(SparkNumericType),
+            SparkArrayType(SparkBinaryType)))),
+      varArgs = true),
     createFunctionWithInputTypes("concat_ws", Seq(SparkStringType, SparkStringType)),
     createFunctionWithInputTypes("contains", Seq(SparkStringType, SparkStringType)),
     createFunctionWithInputTypes("ends_with", Seq(SparkStringType, SparkStringType)),

fuzz-testing/src/main/scala/org/apache/comet/fuzz/QueryGen.scala

@@ -21,6 +21,7 @@ package org.apache.comet.fuzz
 
 import java.io.{BufferedWriter, FileWriter}
 
+import scala.annotation.tailrec
 import scala.collection.mutable
 import scala.util.Random
 
@@ -103,7 +104,12 @@ object QueryGen {
     val func = Utils.randomChoice(Meta.scalarFunc, r)
     try {
       val signature = Utils.randomChoice(func.signatures, r)
-      val args = signature.inputTypes.map(x => pickRandomColumn(r, table, x))
+      val args =
+        if (signature.varArgs) {
+          pickRandomColumns(r, table, signature.inputTypes.head)
+        } else {
+          signature.inputTypes.map(x => pickRandomColumn(r, table, x))
+        }
 
       // Example SELECT c0, log(c0) as x FROM test0
       s"SELECT ${args.mkString(", ")}, ${func.name}(${args.mkString(", ")}) AS x " +
@@ -117,6 +123,21 @@ object QueryGen {
     }
   }
 
+  @tailrec
+  private def pickRandomColumns(r: Random, df: DataFrame, targetType: SparkType): Seq[String] = {
+    targetType match {
+      case SparkTypeOneOf(choices) =>
+        val chosenType = Utils.randomChoice(choices, r)
+        pickRandomColumns(r, df, chosenType)
+      case _ =>
+        var columns = Set.empty[String]
+        for (_ <- 0 to r.nextInt(df.columns.length)) {
+          columns += pickRandomColumn(r, df, targetType)
+        }
+        columns.toSeq
+    }
+  }
+
   private def pickRandomColumn(r: Random, df: DataFrame, targetType: SparkType): String = {
     targetType match {
       case SparkAnyType =>