[SPARK-23736][SQL] Extending the concat function to support array columns

python/pyspark/sql/functions.py

@@ -1414,21 +1414,6 @@ def hash(*cols):
 del _name, _doc
 
 
-@since(1.5)
-@ignore_unicode_prefix
-def concat(*cols):
-    """
-    Concatenates multiple input columns together into a single column.
-    If all inputs are binary, concat returns an output as binary. Otherwise, it returns as string.
-
-    >>> df = spark.createDataFrame([('abcd','123')], ['s', 'd'])
-    >>> df.select(concat(df.s, df.d).alias('s')).collect()
-    [Row(s=u'abcd123')]
-    """
-    sc = SparkContext._active_spark_context
-    return Column(sc._jvm.functions.concat(_to_seq(sc, cols, _to_java_column)))
-
-
 @since(1.5)
 @ignore_unicode_prefix
 def concat_ws(sep, *cols):
@@ -1834,6 +1819,25 @@ def array_contains(col, value):
     return Column(sc._jvm.functions.array_contains(_to_java_column(col), value))
 
 
+@since(1.5)
+@ignore_unicode_prefix
+def concat(*cols):
+    """
+    Concatenates multiple input columns together into a single column.
+    The function works with strings, binary columns and arrays of the same time.
+
+    >>> df = spark.createDataFrame([('abcd','123')], ['s', 'd'])
+    >>> df.select(concat(df.s, df.d).alias('s')).collect()
+    [Row(s=u'abcd123')]
+
+    >>> df = spark.createDataFrame([([1, 2], [3, 4], [5]), ([1, 2], None, [3])], ['a', 'b', 'c'])
+    >>> df.select(concat(df.a, df.b, df.c).alias("arr")).collect()
+    [Row(arr=[1, 2, 3, 4, 5]), Row(arr=None)]
+    """
+    sc = SparkContext._active_spark_context
+    return Column(sc._jvm.functions.concat(_to_seq(sc, cols, _to_java_column)))
+
+
 @since(1.4)
 def explode(col):
     """Returns a new row for each element in the given array or map.

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala

@@ -825,6 +825,8 @@ class Analyzer(
         result
       case UnresolvedExtractValue(child, fieldExpr) if child.resolved =>
         ExtractValue(child, fieldExpr, resolver)
+      case UnresolvedConcat(children) if children.forall(_.resolved) =>
+        ResolveConcat(children)
       case _ => e.mapChildren(resolve(_, q))
     }
 

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala

@@ -308,7 +308,6 @@ object FunctionRegistry {
     expression[BitLength]("bit_length"),
     expression[Length]("char_length"),
     expression[Length]("character_length"),
-    expression[Concat]("concat"),
     expression[ConcatWs]("concat_ws"),
     expression[Decode]("decode"),
     expression[Elt]("elt"),
@@ -408,6 +407,7 @@ object FunctionRegistry {
     expression[MapValues]("map_values"),
     expression[Size]("size"),
     expression[SortArray]("sort_array"),
+    expression[UnresolvedConcat]("concat"),
     CreateStruct.registryEntry,
 
     // misc functions

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/unresolved.scala

@@ -480,3 +480,24 @@ case class UnresolvedOrdinal(ordinal: Int)
   override def nullable: Boolean = throw new UnresolvedException(this, "nullable")
   override lazy val resolved = false
 }
+
+/**
+ * Concatenates multiple columns of the same type into one.
+ * @param children Could be string, binary or array expressions
+ */
+@ExpressionDescription(
+  usage = "_FUNC_(col1, col2, ..., colN) - Returns the concatenation of col1, col2, ..., colN.",
+  examples = """
+    Examples:
+      > SELECT _FUNC_('Spark', 'SQL');
+       SparkSQL
+      > SELECT _FUNC_(array(1, 2, 3), array(4, 5), array(6));
+       [1,2,3,4,5,6]
+  """)
+case class UnresolvedConcat(children: Seq[Expression]) extends Expression
+  with Unevaluable {
+  override def dataType: DataType = throw new UnresolvedException(this, "dataType")
+  override def foldable: Boolean = throw new UnresolvedException(this, "foldable")
+  override def nullable: Boolean = throw new UnresolvedException(this, "nullable")
+  override lazy val resolved = false
+}

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Expression.scala

@@ -699,3 +699,90 @@ abstract class TernaryExpression extends Expression {
  * and Hive function wrappers.
  */
 trait UserDefinedExpression
+
+/**
+ * The trait covers logic for performing null safe evaluation and code generation.
+ */
+trait NullSafeEvaluation extends Expression
+{
+  override def foldable: Boolean = children.forall(_.foldable)
+
+  override def nullable: Boolean = children.exists(_.nullable)
+
+  /**
+   * Default behavior of evaluation according to the default nullability of NullSafeEvaluation.
+   * If a class utilizing NullSaveEvaluation override [[nullable]], probably should also
+   * override this.
+   */
+  override def eval(input: InternalRow): Any =
+  {
+    val values = children.toStream.map(_.eval(input))
+    if (values.contains(null)) {
+      null
+    } else {
+      nullSafeEval(values)
+    }
+  }
+
+  /**
+   * Called by default [[eval]] implementation. If a class utilizing NullSaveEvaluation keep
+   * the default nullability, they can override this method to save null-check code.  If we need
+   * full control of evaluation process, we should override [[eval]].
+   */
+  protected def nullSafeEval(inputs: Seq[Any]): Any =
+    sys.error(s"The class utilizing NullSaveEvaluation must override either eval or nullSafeEval")
+
+  /**
+   * Short hand for generating of null save evaluation code.
+   * If either of the sub-expressions is null, the result of this computation
+   * is assumed to be null.
+   *
+   * @param f accepts a sequence of variable names and returns Java code to compute the output.
+   */
+  protected def defineCodeGen(
+    ctx: CodegenContext,
+    ev: ExprCode,
+    f: Seq[String] => String): ExprCode = {
+    nullSafeCodeGen(ctx, ev, values => {
+      s"${ev.value} = ${f(values)};"
+    })
+  }
+
+  /**
+   * Called by expressions to generate null safe evaluation code.
+   * If either of the sub-expressions is null, the result of this computation
+   * is assumed to be null.
+   *
+   * @param f a function that accepts a sequence of non-null evaluation result names of children
+   *          and returns Java code to compute the output.
+   */
+  protected def nullSafeCodeGen(
+   ctx: CodegenContext,
+   ev: ExprCode,
+   f: Seq[String] => String): ExprCode = {
+    val gens = children.map(_.genCode(ctx))
+    val resultCode = f(gens.map(_.value))
+
+    if (nullable) {
+      val nullSafeEval = children.zip(gens).foldRight(s"""
+          ${ev.isNull} = false; // resultCode could change nullability.
+          $resultCode
+        """) {
+          case ((child, gen), acc) =>
+            gen.code + ctx.nullSafeExec(child.nullable, gen.isNull)(acc)
+        }
+
+      ev.copy(code = s"""
+        boolean ${ev.isNull} = true;
+        ${CodeGenerator.javaType(dataType)} ${ev.value} = ${CodeGenerator.defaultValue(dataType)};
+        $nullSafeEval
+      """)
+    } else {
+      ev.copy(code = s"""
+        boolean ${ev.isNull} = false;
+        ${gens.map(_.code).mkString("\n")}
+        ${CodeGenerator.javaType(dataType)} ${ev.value} = ${CodeGenerator.defaultValue(dataType)};
+        $resultCode""", isNull = "false")
+    }
+  }
+}

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collectionOperations.scala

@@ -21,8 +21,10 @@ import java.util.Comparator
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.analysis.TypeCheckResult
 import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext, CodeGenerator, CodegenFallback, ExprCode}
-import org.apache.spark.sql.catalyst.util.{ArrayData, GenericArrayData, MapData}
+import org.apache.spark.sql.catalyst.util.{ArrayData, GenericArrayData, MapData, TypeUtils}
 import org.apache.spark.sql.types._
+import org.apache.spark.unsafe.Platform
+import org.apache.spark.unsafe.array.ByteArrayMethods
 
 /**
  * Given an array or map, returns its size. Returns -1 if null.
@@ -287,3 +289,171 @@ case class ArrayContains(left: Expression, right: Expression)
 
   override def prettyName: String = "array_contains"
 }
+
+/**
+ * Replaces [[org.apache.spark.sql.catalyst.analysis.UnresolvedConcat UnresolvedConcat]]s
+ * with concrete concate expressions.
+ */
+object ResolveConcat
+{
+  def apply(children: Seq[Expression]): Expression = {
+    if (children.nonEmpty && ArrayType.acceptsType(children(0).dataType)) {
+      ConcatArrays(children)
+    } else {
+      Concat(children)
+    }
+  }
+}
+
+/**
+ * Concatenates multiple arrays into one.
+ */
+@ExpressionDescription(
+  usage = "_FUNC_(expr, ...) - Concatenates multiple arrays of the same type into one.",
+  examples = """
+    Examples:
+      > SELECT _FUNC_(array(1, 2, 3), array(4, 5), array(6));
+       [1,2,3,4,5,6]
+  """,
+  since = "2.4.0")
+case class ConcatArrays(children: Seq[Expression]) extends Expression with NullSafeEvaluation {
+
+  override def checkInputDataTypes(): TypeCheckResult = {
+    val arrayCheck = checkInputDataTypesAreArrays
+    if(arrayCheck.isFailure) {
+      arrayCheck
+    } else {
+      TypeUtils.checkForSameTypeInputExpr(children.map(_.dataType), s"function $prettyName")
+    }
+  }
+
+  private def checkInputDataTypesAreArrays(): TypeCheckResult =
+  {
+    val mismatches = children.zipWithIndex.collect {
+      case (child, idx) if !ArrayType.acceptsType(child.dataType) =>
+        s"argument ${idx + 1} has to be ${ArrayType.simpleString} type, " +
+          s"however, '${child.sql}' is of ${child.dataType.simpleString} type."
+    }
+
+    if (mismatches.isEmpty) {
+      TypeCheckResult.TypeCheckSuccess
+    } else {
+      TypeCheckResult.TypeCheckFailure(mismatches.mkString(" "))
+    }
+  }
+
+  override def dataType: ArrayType =
+    children
+      .headOption.map(_.dataType.asInstanceOf[ArrayType])
+      .getOrElse(ArrayType.defaultConcreteType.asInstanceOf[ArrayType])
+
+
+  override protected def nullSafeEval(inputs: Seq[Any]): Any = {
+    val elements = inputs.flatMap(_.asInstanceOf[ArrayData].toObjectArray(dataType.elementType))
+    new GenericArrayData(elements)
+  }
+
+  override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+    nullSafeCodeGen(ctx, ev, arrays => {
+      val elementType = dataType.elementType
+      if (CodeGenerator.isPrimitiveType(elementType)) {
+        genCodeForConcatOfPrimitiveElements(ctx, elementType, arrays, ev.value)
+      } else {
+        genCodeForConcatOfComplexElements(ctx, arrays, ev.value)
+      }
+    })
+  }
+
+  private def genCodeForNumberOfElements(
+    ctx: CodegenContext,
+    elements: Seq[String]
+  ) : (String, String) = {
+    val variableName = ctx.freshName("numElements")
+    val code = elements
+      .map(el => s"$variableName += $el.numElements();")
+      .foldLeft(s"int $variableName = 0;")((acc, s) => acc + "\n" + s)
+    (code, variableName)
+  }
+
+  private def genCodeForConcatOfPrimitiveElements(
+    ctx: CodegenContext,
+    elementType: DataType,
+    elements: Seq[String],
+    arrayDataName: String
+  ): String = {
+    val arrayName = ctx.freshName("array")
+    val arraySizeName = ctx.freshName("size")
+    val counter = ctx.freshName("counter")
+    val tempArrayDataName = ctx.freshName("tempArrayData")
+
+    val (numElemCode, numElemName) = genCodeForNumberOfElements(ctx, elements)
+
+    val unsafeArraySizeInBytes = s"""
+      |int $arraySizeName = UnsafeArrayData.calculateHeaderPortionInBytes($numElemName) +
+      |${classOf[ByteArrayMethods].getName}.roundNumberOfBytesToNearestWord(
+      |  ${elementType.defaultSize} * $numElemName
+      |);
+      """.stripMargin
+    val baseOffset = Platform.BYTE_ARRAY_OFFSET
+
+    val primitiveValueTypeName = CodeGenerator.primitiveTypeName(elementType)
+    val assignments = elements.map { el =>
+      s"""
+        |for (int z = 0; z < $el.numElements(); z++) {
+        |  if ($el.isNullAt(z)) {
+        |    $tempArrayDataName.setNullAt($counter);
+        |  } else {
+        |    $tempArrayDataName.set$primitiveValueTypeName(
+        |      $counter,
+        |      $el.get$primitiveValueTypeName(z)
+        |    );
+        |  }
+        |  $counter++;
+        |}
+        """.stripMargin
+    }.mkString("\n")
+
+    s"""
+      |$numElemCode
+      |$unsafeArraySizeInBytes
+      |byte[] $arrayName = new byte[$arraySizeName];
+      |UnsafeArrayData $tempArrayDataName = new UnsafeArrayData();
+      |Platform.putLong($arrayName, $baseOffset, $numElemName);
+      |$tempArrayDataName.pointTo($arrayName, $baseOffset, $arraySizeName);
+      |int $counter = 0;
+      |$assignments
+      |$arrayDataName = $tempArrayDataName;
+      """.stripMargin
+
+  }
+
+  private def genCodeForConcatOfComplexElements(
+   ctx: CodegenContext,
+   elements: Seq[String],
+   arrayDataName: String
+  ): String = {
+    val genericArrayClass = classOf[GenericArrayData].getName
+    val arrayName = ctx.freshName("arrayObject")
+    val counter = ctx.freshName("counter")
+    val (numElemCode, numElemName) = genCodeForNumberOfElements(ctx, elements)
+
+    val assignments = elements.map { el =>
+      s"""
+        |for (int z = 0; z < $el.numElements(); z++) {
+        |  $arrayName[$counter] = $el.array()[z];
+        |  $counter++;
+        |}
+        """.stripMargin
+    }.mkString("\n")
+
+    s"""
+      |$numElemCode
+      |Object[] $arrayName = new Object[$numElemName];
+      |int $counter = 0;
+      |$assignments
+      |$arrayDataName = new $genericArrayClass($arrayName);
+      """.stripMargin
+  }
+
+  override def prettyName: String = "concat"
+}