uses splitexpression for inputs

Signed-off-by: DylanGuedes <[email protected]>

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

@@ -184,18 +184,25 @@ case class Zip(children: Seq[Expression]) extends Expression with ExpectsInputTy
         |  $arrCardinality[$index] = 0;
         |}
         |$storedArrTypes[$index] = "${arrayElementTypes(index)}";
-        |$biggestCardinality[0] = Math.max($biggestCardinality[0], $arrCardinality[$index]);
+        |$biggestCardinality = Math.max($biggestCardinality, $arrCardinality[$index]);
       """.stripMargin
     }
 
     val inputsSplitted = ctx.splitExpressions(
       expressions = inputs,
       funcName = "getInputAndCardinality",
+      returnType = "int",
+      makeSplitFunction = body =>
+        s"""
+          |$body
+          |return $biggestCardinality;
+        """.stripMargin,
+      foldFunctions = _.map(funcCall => s"$biggestCardinality = $funcCall;").mkString("\n"),
       arguments =
         ("ArrayData[]", arrVals) ::
         ("int[]", arrCardinality) ::
         ("String[]", storedArrTypes) ::
-        ("int[]", biggestCardinality) :: Nil)
+        ("int", biggestCardinality) :: Nil)
 
     val myobject = ctx.freshName("myobject")
     val j = ctx.freshName("j")
@@ -205,32 +212,24 @@ case class Zip(children: Seq[Expression]) extends Expression with ExpectsInputTy
     val fillValue = arrayElementTypes.distinct.map { case (elementType) =>
       val getArrValsItem = CodeGenerator.getValue(s"$arrVals[$j]", elementType, i)
       s"""
-      |      if ($storedArrTypes[$j] == "${elementType}") {
-      |        $myobject[$j] = $getArrValsItem;
-      |      }
+      |if ($storedArrTypes[$j] == "${elementType}") {
+      |  $myobject[$j] = $getArrValsItem;
+      |}
       """.stripMargin
-    }
-
-    val fillValueSplitted = ctx.splitExpressions(
-      expressions = fillValue,
-      funcName = "fillValue",
-      arguments =
-        ("int", j) ::
-        ("Array[] Object", myobject) :: Nil)
+    }.mkString("\n")
 
     ev.copy(s"""
       |ArrayData[] $arrVals = new ArrayData[$numberOfArrays];
       |int[] $arrCardinality = new int[$numberOfArrays];
-      |int[] $biggestCardinality = new int[1];
-      |$biggestCardinality[0] = 0;
+      |int $biggestCardinality = 0;
       |String[] $storedArrTypes = new String[$numberOfArrays];
       |$inputsSplitted
-      |Object[] $args = new Object[$biggestCardinality[0]];
-      |for (int $i = 0; $i < $biggestCardinality[0]; $i ++) {
+      |Object[] $args = new Object[$biggestCardinality];
+      |for (int $i = 0; $i < $biggestCardinality; $i ++) {
       |  Object[] $myobject = new Object[$numberOfArrays];
       |  for (int $j = 0; $j < $numberOfArrays; $j ++) {
       |    if ($arrVals[$j] != null && $arrCardinality[$j] > $i && !$arrVals[$j].isNullAt($i)) {
-      |      $fillValueSplitted
+      |      $fillValue
       |    } else {
       |      $myobject[$j] = null;
       |    }

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

@@ -18,8 +18,8 @@
 package org.apache.spark.sql.catalyst.expressions
 
 import org.apache.spark.SparkFunSuite
-import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.Row
+import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.types._
 
 class CollectionExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper {
@@ -377,6 +377,25 @@ class CollectionExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper
 
     checkEvaluation(Zip(Seq(literals(7), literals(10))),
       List(Row(null, Array[Byte](1.toByte, 5.toByte))))
+
+    val longLiteral =
+      Literal.create((0 to 1000).toSeq, ArrayType(IntegerType))
+
+    checkEvaluation(Zip(Seq(literals(0), longLiteral)),
+      List(Row(9001, 0), Row(9002, 1), Row(9003, 2)) ++
+      (3 to 1000).map { Row(null, _) }.toList)
+
+    val manyLiterals = (0 to 1000).map { case (number) =>
+      Literal.create(Seq(1), ArrayType(IntegerType))
+    }.toSeq
+
+    val numbers = List(
+      Seq(9001) ++ (0 to 1000).map { case (number) => 1 }.toSeq,
+      Seq(9002) ++ (0 to 1000).map { case (number) => null }.toSeq,
+      Seq(9003) ++ (0 to 1000).map { case (number) => null }.toSeq,
+      Seq(null) ++ (0 to 1000).map { case (number) => null }.toSeq)
+    checkEvaluation(Zip(Seq(literals(0)) ++ manyLiterals),
+      List(Row(numbers(0): _*), Row(numbers(1): _*), Row(numbers(2): _*), Row(numbers(3): _*)))
   }
 
   test("Array Min") {