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

python/pyspark/sql/functions.py

@@ -1834,6 +1834,25 @@ def array_contains(col, value):
     return Column(sc._jvm.functions.array_contains(_to_java_column(col), value))
 
 
+@since(2.4)
+def concat_arrays(*cols):
+    """
+    Collection function: Concatenates multiple arrays into one.
+
+    :param cols: list of column names (string) or list of :class:`Column` expressions that have
+        the same data type.
+
+    >>> df = spark.createDataFrame([([1, 2], [3, 4], [5]), ([1, 2], None, [3])], ['a', 'b', 'c'])
+    >>> df.select(concat_arrays(df.a, df.b, df.c).alias("arr")).collect()
+    [Row(arr=[1, 2, 3, 4, 5]), Row(arr=None)]
+    """
+    sc = SparkContext._active_spark_context
+    if len(cols) == 1 and isinstance(cols[0], (list, set)):
+        cols = cols[0]
+    args = _to_seq(sc, cols, _to_java_column)
+    return Column(sc._jvm.functions.concat_arrays(args))
+
+
 @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/FunctionRegistry.scala

@@ -408,6 +408,7 @@ object FunctionRegistry {
     expression[MapValues]("map_values"),
     expression[Size]("size"),
     expression[SortArray]("sort_array"),
+    expression[ConcatArrays]("concat_arrays"),
     CreateStruct.registryEntry,
 
     // misc functions

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

@@ -699,3 +699,88 @@ abstract class TernaryExpression extends Expression {
  * and Hive function wrappers.
  */
 trait UserDefinedExpression
+
+/**
+ * The trait covers logic for performing null save 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.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 =
+        (s"""
+          ${ev.isNull} = false; // resultCode could change nullability.
+          $resultCode
+        """ /: children.zip(gens)) {
+          case (acc, (child, gen)) =>
+            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,152 @@ case class ArrayContains(left: Expression, right: Expression)
 
   override def prettyName: String = "array_contains"
 }
+
+/**
+ * Concatenates multiple arrays into one.
+ */
+@ExpressionDescription(
+  usage = "_FUNC_(expr, ...) - Concatenates multiple arrays into one.",
+  examples = """
+    Examples:
+      > SELECT _FUNC_(array(1, 2, 3), array(4, 5), array(6));
+       [1,2,3,4,5,6]
+  """)
+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_arrays"
+}

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CollectionExpressionsSuite.scala

@@ -105,4 +105,30 @@ class CollectionExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper
     checkEvaluation(ArrayContains(a3, Literal("")), null)
     checkEvaluation(ArrayContains(a3, Literal.create(null, StringType)), null)
   }
+
+  test("Concat Arrays") {
+    val a0 = Literal.create(Seq(1, 2, 3), ArrayType(IntegerType))
+    val a1 = Literal.create(Seq.empty[Integer], ArrayType(IntegerType))
+    val a2 = Literal.create(Seq(4, null, 6), ArrayType(IntegerType))
+    val a3 = Literal.create(Seq("a", "b", "c"), ArrayType(StringType))
+    val a4 = Literal.create(Seq("e", null), ArrayType(StringType))
+    val a5 = Literal.create(Seq.empty[String], ArrayType(StringType))
+    val an = Literal.create(null, ArrayType(StringType))
+
+    checkEvaluation(ConcatArrays(Seq(a0)), Seq(1, 2, 3))
+    checkEvaluation(ConcatArrays(Seq(a0, a2)), Seq(1, 2, 3, 4, null, 6))
+    checkEvaluation(ConcatArrays(Seq(a0, a1, a2)), Seq(1, 2, 3, 4, null, 6))
+
+    checkEvaluation(ConcatArrays(Seq(a3, a4)), Seq("a", "b", "c", "e", null))
+    checkEvaluation(ConcatArrays(Seq(a3, a4, a5)), Seq("a", "b", "c", "e", null))
+    checkEvaluation(ConcatArrays(Seq(a5)), Seq())
+
+    checkEvaluation(ConcatArrays(Seq()), Seq())
+    checkEvaluation(ConcatArrays(Seq(a0, a0)), Seq(1, 2, 3, 1, 2, 3))
+
+    checkEvaluation(ConcatArrays(Seq(an)), null)
+    checkEvaluation(ConcatArrays(Seq(a3, an)), null)
+    checkEvaluation(ConcatArrays(Seq(an, a3)), null)
+    checkEvaluation(ConcatArrays(Seq(a3, an, a4)), null)
+  }
 }

sql/core/src/main/scala/org/apache/spark/sql/functions.scala

@@ -3046,6 +3046,14 @@ object functions {
     ArrayContains(column.expr, Literal(value))
   }
 
+  /**
+   * Merges multiple arrays into one by putting elements from the specific array after elements
+   * from the previous array. If any of the arrays is null, null is returned.
+   * @group collection_funcs
+   * @since 2.4.0
+   */
+  def concat_arrays(columns: Column*): Column = withExpr { ConcatArrays(columns.map(_.expr)) }
+
   /**
    * Creates a new row for each element in the given array or map column.
    *

sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala

@@ -413,6 +413,63 @@ class DataFrameFunctionsSuite extends QueryTest with SharedSQLContext {
     )
   }
 
+  test("concat arrays function") {
+    val nint : Int = null.asInstanceOf[Int]
+    val nseqi : Seq[Int] = null
+    val nseqs : Seq[String] = null
+    val df = Seq(
+      (Seq(1), Seq(2, 3, 4), Seq(5, 6), nseqi, Seq("a", "b", "c"), Seq("d", "e"), Seq("f"), nseqs),
+      (Seq(1, nint), Seq.empty[Int], Seq(2), nseqi, Seq("a"), Seq.empty[String], Seq(null), nseqs)
+    ).toDF("i1", "i2", "i3", "in", "s1", "s2", "s3", "sn")
+
+    // Simple test cases
+    checkAnswer(
+      df.select(concat_arrays($"i1", $"i2", $"i3")),
+      Seq(Row(Seq(1, 2, 3, 4, 5, 6)), Row(Seq(1, nint, 2)))
+    )
+    checkAnswer(
+      df.selectExpr("concat_arrays(i1, i2, i3)"),
+      Seq(Row(Seq(1, 2, 3, 4, 5, 6)), Row(Seq(1, nint, 2)))
+    )
+    checkAnswer(
+      df.select(concat_arrays($"s1", $"s2", $"s3")),
+      Seq(Row(Seq("a", "b", "c", "d", "e", "f")), Row(Seq("a", null)))
+    )
+    checkAnswer(
+      df.selectExpr("concat_arrays(s1, s2, s3)"),
+      Seq(Row(Seq("a", "b", "c", "d", "e", "f")), Row(Seq("a", null)))
+    )
+
+    // Null test cases
+    checkAnswer(
+      df.select(concat_arrays($"i1", $"in")),
+      Seq(Row(null), Row(null))
+    )
+    checkAnswer(
+      df.select(concat_arrays($"in", $"i1")),
+      Seq(Row(null), Row(null))
+    )
+    checkAnswer(
+      df.select(concat_arrays($"s1", $"sn")),
+      Seq(Row(null), Row(null))
+    )
+    checkAnswer(
+      df.select(concat_arrays($"sn", $"s1")),
+      Seq(Row(null), Row(null))
+    )
+
+    // Type error test cases
+    intercept[AnalysisException] {
+      df.select(concat_arrays($"i1", $"s1"))
+    }
+    intercept[AnalysisException] {
+      df.select(concat_arrays(lit("a"), lit("b")))
+    }
+    intercept[AnalysisException] {
+      df.selectExpr("concat_arrays(i1, i2, null)")
+    }
+  }
+
   private def assertValuesDoNotChangeAfterCoalesceOrUnion(v: Column): Unit = {
     import DataFrameFunctionsSuite.CodegenFallbackExpr
     for ((codegenFallback, wholeStage) <- Seq((true, false), (false, false), (false, true))) {