infoankitp commented on code in PR #38865:
URL: https://github.com/apache/spark/pull/38865#discussion_r1061145796
##########
sql/core/src/test/resources/sql-tests/inputs/array.sql:
##########
@@ -137,3 +137,11 @@ select array_compact(array("a", "b", "c"));
select array_compact(array(1D, null, 2D, null));
select array_compact(array(array(1, 2, 3, null), null, array(4, null, 6)));
select array_compact(array(null));
+
+-- function array_append
Review Comment:
Added the tests but had to add the type cast for empty array in second one.
Else it was returning null
select array_append(CAST(array() AS ARRAY<String>), CAST(NULL AS String));
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collectionOperations.scala:
##########
@@ -4630,4 +4630,144 @@ case class ArrayCompact(child: Expression)
copy(child = newChild)
}
+/**
+ * Given an array, and another element append the element at the end of the
array.
+ */
+@ExpressionDescription(
+ usage = """
+ _FUNC_(array, element) - Add the element at the end of the array passed
as first
+ argument. Type of element should be similar to type of the elements of
the array.
+ """,
+ examples = """
+ Examples:
+ > SELECT _FUNC_(array('b', 'd', 'c', 'a'), 'd');
+ ["b","d","c","a","d"]
+ """,
+ since = "3.4.0",
+ group = "array_funcs")
+case class ArrayAppend(left: Expression, right: Expression)
+ extends BinaryExpression
+ with ImplicitCastInputTypes
+ with ComplexTypeMergingExpression
+ with QueryErrorsBase {
+ override def prettyName: String = "array_append"
+
+ @transient protected lazy val elementType: DataType =
+ inputTypes.head.asInstanceOf[ArrayType].elementType
+
+ override def inputTypes: Seq[AbstractDataType] = {
+ (left.dataType, right.dataType) match {
+ case (ArrayType(e1, hasNull), e2) =>
+ TypeCoercion.findTightestCommonType(e1, e2) match {
+ case Some(dt) => Seq(ArrayType(dt, hasNull), dt)
+ case _ => Seq.empty
+ }
+ case _ => Seq.empty
+ }
+ }
+
+ override def checkInputDataTypes(): TypeCheckResult = {
+ (left.dataType, right.dataType) match {
+ case (ArrayType(e1, _), e2) if e1.sameType(e2) =>
TypeCheckResult.TypeCheckSuccess
+ case (ArrayType(e1, _), e2) => DataTypeMismatch(
+ errorSubClass = "ARRAY_FUNCTION_DIFF_TYPES",
+ messageParameters = Map(
+ "functionName" -> toSQLId(prettyName),
+ "leftType" -> toSQLType(left.dataType),
+ "rightType" -> toSQLType(right.dataType),
+ "dataType" -> toSQLType(ArrayType)
+ ))
+ case _ =>
+ DataTypeMismatch(
+ errorSubClass = "UNEXPECTED_INPUT_TYPE",
+ messageParameters = Map(
+ "paramIndex" -> "0",
+ "requiredType" -> toSQLType(ArrayType),
+ "inputSql" -> toSQLExpr(left),
+ "inputType" -> toSQLType(left.dataType)
+ )
+ )
+ }
+ }
+
+ override def eval(input: InternalRow): Any = {
+ val value1 = left.eval(input)
+ if (value1 == null) {
+ null
+ } else {
+ val value2 = right.eval(input)
+ nullSafeEval(value1, value2)
+ }
+ }
+
+ override protected def nullSafeEval(arr: Any, elementData: Any): Any = {
+ val arrayData = arr.asInstanceOf[ArrayData]
+ val numberOfElements = arrayData.numElements() + 1
+ if (numberOfElements > ByteArrayMethods.MAX_ROUNDED_ARRAY_LENGTH) {
+ throw
QueryExecutionErrors.concatArraysWithElementsExceedLimitError(numberOfElements)
+ }
+ val finalData = new Array[Any](numberOfElements)
+ arrayData.foreach(elementType, finalData.update)
+ finalData.update(arrayData.numElements(), elementData)
+ new GenericArrayData(finalData)
+ }
+
+ override def nullable: Boolean = left.nullable
+ override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+ val leftGen = left.genCode(ctx)
+ val rightGen = right.genCode(ctx)
+ val f = (eval1: String, eval2: String) => {
+ val newArraySize = ctx.freshName("newArraySize")
+ val i = ctx.freshName("i")
+ val values = ctx.freshName("values")
+ val allocation = CodeGenerator.createArrayData(
+ values, elementType, newArraySize, s" $prettyName failed.")
+ val assignment = CodeGenerator.createArrayAssignment(
+ values, elementType, eval1, i, i,
left.dataType.asInstanceOf[ArrayType].containsNull)
+ s"""
+ |int $newArraySize = $eval1.numElements() + 1;
+ |$allocation
+ |int $i = 0;
+ |while ($i < $eval1.numElements()) {
+ | $assignment
+ | $i ++;
+ |}
+ |${CodeGenerator.setArrayElement(values, elementType, i, eval2,
Some(rightGen.isNull))}
+ |${ev.value} = $values;
+ |""".stripMargin
+ }
+ val resultCode = f(leftGen.value, rightGen.value)
+ if (nullable) {
+ val nullSafeEval =
+ leftGen.code + rightGen.code + ctx.nullSafeExec(left.nullable,
leftGen.isNull) {
+ s"""
+ ${ev.isNull} = false; // resultCode could change nullability.
+ $resultCode
+ """
+ }
+ ev.copy(code =
+ code"""
+ boolean ${ev.isNull} = true;
+ ${CodeGenerator.javaType(dataType)} ${ev.value} =
${CodeGenerator.defaultValue(dataType)};
+ $nullSafeEval
+ """)
+ }
+ else {
Review Comment:
Done
--
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.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
