feat: support array_repeat

Author: comphead

spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala

@@ -1979,6 +1979,7 @@ object QueryPlanSerde extends Logging with CometExprShim {
       case _: ArrayIntersect => convert(CometArrayIntersect)
       case _: ArrayJoin => convert(CometArrayJoin)
       case _: ArraysOverlap => convert(CometArraysOverlap)
+      case _: ArrayRepeat => convert(CometArrayRepeat)
       case _ @ArrayFilter(_, func) if func.children.head.isInstanceOf[IsNotNull] =>
         convert(CometArrayCompact)
       case _: ArrayExcept =>
@@ -3068,7 +3069,7 @@ trait CometAggregateExpressionSerde {
    * Convert a Spark expression into a protocol buffer representation that can be passed into
    * native code.
    *
-   * @param expr
+   * @param aggExpr
    *   The aggregate expression.
    * @param expr
    *   The aggregate function.

spark/src/main/scala/org/apache/comet/serde/arrays.scala

@@ -179,6 +179,20 @@ object CometArraysOverlap extends CometExpressionSerde with IncompatExpr {
   }
 }
 
+object CometArrayRepeat extends CometExpressionSerde with IncompatExpr {
+  override def convert(
+      expr: Expression,
+      inputs: Seq[Attribute],
+      binding: Boolean): Option[ExprOuterClass.Expr] = {
+    val leftArrayExprProto = exprToProto(expr.children.head, inputs, binding)
+    val rightArrayExprProto = exprToProto(expr.children(1), inputs, binding)
+
+    val arraysRepeatScalarExpr =
+      scalarExprToProto("array_repeat", leftArrayExprProto, rightArrayExprProto)
+    optExprWithInfo(arraysRepeatScalarExpr, expr, expr.children: _*)
+  }
+}
+
 object CometArrayCompact extends CometExpressionSerde with IncompatExpr {
   override def convert(
       expr: Expression,

spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala

@@ -410,4 +410,26 @@ class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelp
     }
   }
 
+  test("array_repeat") {
+    withSQLConf(
+      CometConf.COMET_EXPR_ALLOW_INCOMPATIBLE.key -> "true",
+      CometConf.COMET_EXPLAIN_FALLBACK_ENABLED.key -> "true") {
+      Seq(true, false).foreach { dictionaryEnabled =>
+        withTempDir { dir =>
+          val path = new Path(dir.toURI.toString, "test.parquet")
+          makeParquetFileAllTypes(path, dictionaryEnabled, 10000)
+          spark.read.parquet(path.toString).createOrReplaceTempView("t1")
+          spark.sql("select * from t1").printSchema()
+
+          checkSparkAnswerAndOperator(sql("SELECT array_repeat(_4, null) from t1"))
+//          checkSparkAnswerAndOperator(
+//            sql("SELECT array_repeat(_2, 5) from t1 where _2 is not null"))
+//          checkSparkAnswerAndOperator(sql("SELECT array_repeat(_3, 2) from t1 where _3 is null"))
+//          checkSparkAnswerAndOperator(sql("SELECT array_repeat(_3, _3) from t1 where _3 is null"))
+//          checkSparkAnswerAndOperator(sql("SELECT array_repeat(cast(_3 as string), 2) from t1"))
+//          checkSparkAnswerAndOperator(sql("SELECT array_repeat(array(_2, _3, _4), 2) from t1"))
+        }
+      }
+    }
+  }
 }