[GitHub] [spark] yaooqinn commented on a change in pull request #26371: [SPARK-27976][SQL] Add built-in Array Functions: array_append

[GitBox <gi...@apache.org>]

yaooqinn commented on a change in pull request #26371: [SPARK-27976][SQL] Add built-in Array Functions: array_append
URL: https://github.com/apache/spark/pull/26371#discussion_r341832868
 
 

 ##########
 File path: sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collectionOperations.scala
 ##########
 @@ -3928,3 +3928,88 @@ case class ArrayExcept(left: Expression, right: Expression) extends ArrayBinaryL
 
   override def prettyName: String = "array_except"
 }
+
+@ExpressionDescription(
+  usage = """
+  _FUNC_(array, element) - Returns an array of appending an element to the end of an array
+  """,
+  examples = """
+    Examples:
+      > SELECT _FUNC_(array(1, 2, 3), 3);
+       [1,2,3,3]
+  """,
+  since = "3.0.0")
+case class ArrayAppend(left: Expression, right: Expression) extends BinaryExpression {
+
+  override def checkInputDataTypes(): TypeCheckResult = left.dataType match {
+    case ArrayType(et, _) => right.dataType match {
+      case NullType =>
+        TypeCheckResult.TypeCheckSuccess
+      case _: AtomicType =>
+        TypeUtils.checkForSameTypeInputExpr(Seq(et, right.dataType), s"function $prettyName")
+      case o => TypeCheckResult.TypeCheckFailure(
+        s"function $prettyName not support append ${o.typeName} to array[${et.typeName}]")
+    }
+    case NullType => right.dataType match {
+      case _: AtomicType =>
+        TypeCheckResult.TypeCheckSuccess
+      case o => TypeCheckResult.TypeCheckFailure(
+        s"function $prettyName not support append ${o.typeName} to array")
+    }
+    case o => TypeCheckResult.TypeCheckFailure(
+      s"function $prettyName not support append to ${o.typeName} type")
+  }
+
+  override def eval(input: InternalRow): Any = {
+    val array = left.eval(input)
+    val elem = right.eval(input)
+    if (array == null) {
+      new GenericArrayData(Array(elem))
+    } else {
+      new GenericArrayData(array.asInstanceOf[ArrayData].array ++ Array(elem))
+    }
+  }
+
+  override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = left.dataType match {
+    case NullType =>
+      val (allocation, assigns, arrayData) =
+        GenArrayData.genCodeToCreateArrayData(ctx, right.dataType, Seq(right), prettyName)
+      ev.copy(
+        code = code"""$allocation$assigns""",
+        value = JavaCode.variable(arrayData, dataType),
+        isNull = FalseLiteral)
+    case ArrayType(et, _) => defineCodeGen(ctx, ev, (inputArray, ne) => {
+      val newArray = ctx.freshName("newArray")
+      val oldArraySize = ctx.freshName("oldArraySize")
+      val newArraySize = ctx.freshName("newArraySize")
+
+      val allocation =
+        CodeGenerator.createArrayData(newArray, et, newArraySize, s"$prettyName failed")
+
+      val i = ctx.freshName("i")
+      val assignment =
+        CodeGenerator.createArrayAssignment(newArray, et, inputArray, i, i, true)
+
+      val setNewValue =
+        CodeGenerator.setArrayElement(newArray, et, oldArraySize, ne, Some(s"$ne == null"))
+
+      s"""
+         |int $oldArraySize = $inputArray.numElements();
+         |int $newArraySize = $inputArray.numElements() + 1;
+         |$allocation
+         |for (int $i = 0; i < $inputArray.numElements(); $i ++) {
+         |  $assignment
+         |}
+         |$setNewValue
+         |${ev.value} = $newArray;
+         |""".stripMargin
+    })
+  }
+
+  override def dataType: DataType = left.dataType match {
 
 Review comment:
   all checked in `checkInputDataTypes`

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
users@infra.apache.org


With regards,
Apache Git Services

---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org