maropu 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_r341829131
########## 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: You need to take care of `right.nullable`. ---------------------------------------------------------------- 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: us...@infra.apache.org With regards, Apache Git Services --------------------------------------------------------------------- To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org For additional commands, e-mail: reviews-h...@spark.apache.org