use zip alias

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

python/pyspark/sql/functions.py

@@ -350,7 +350,7 @@ def corr(col1, col2):
 
     >>> a = range(20)
     >>> b = [2 * x for x in range(20)]
-    >>> df = spark.createDataFrame(__builtin__.zip(a, b), ["a", "b"])
+    >>> df = spark.createDataFrame(zip(a, b), ["a", "b"])
     >>> df.agg(corr("a", "b").alias('c')).collect()
     [Row(c=1.0)]
     """
@@ -364,7 +364,7 @@ def covar_pop(col1, col2):
 
     >>> a = [1] * 10
     >>> b = [1] * 10
-    >>> df = spark.createDataFrame(__builtin__.zip(a, b), ["a", "b"])
+    >>> df = spark.createDataFrame(zip(a, b), ["a", "b"])
     >>> df.agg(covar_pop("a", "b").alias('c')).collect()
     [Row(c=0.0)]
     """
@@ -2402,10 +2402,10 @@ def zip(*cols):
 
     :param cols: columns in input
 
-    >>> from pyspark.sql.functions import zip
+    >>> from pyspark.sql.functions import zip as spark_zip
     >>> df = spark.createDataFrame([(([1, 2, 3], [2, 3, 4]))], ['vals1', 'vals2'])
-    >>> df.select(zip(df.vals1, df.vals2).alias('zipped')).collect()
-    [Row(zipped=[1, 2]), Row(zipped=[2, 3]), Row(zipped=[3, 4])]
+    >>> df.select(spark_zip(df.vals1, df.vals2).alias('zipped')).collect()
+    [Row(zipped=[Row(vals1=1, vals2=2), Row(vals1=2, vals2=3), Row(vals1=3, vals2=4)])]
     """
     sc = SparkContext._active_spark_context
     return Column(sc._jvm.functions.zip(_to_seq(sc, cols, _to_java_column)))

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

@@ -156,30 +156,29 @@ case class Zip(children: Seq[Expression]) extends Expression with ExpectsInputTy
       case ((expr: NamedExpression, elementType), _) =>
         StructField(expr.name, elementType, nullable = true)
       case ((_, elementType), idx) =>
-        StructField(s"$idx", elementType, nullable = true)
+        StructField(idx.toString, elementType, nullable = true)
     }
     StructType(fields)
   }
 
-  val numberOfArrays: Int = children.length
+  @transient lazy val numberOfArrays: Int = children.length
 
-  def emptyInputGenCode(ev: ExprCode): ExprCode = {
-    val genericArrayData = classOf[GenericArrayData].getName
+  @transient lazy val genericArrayData = classOf[GenericArrayData].getName
 
+  def emptyInputGenCode(ev: ExprCode): ExprCode = {
     ev.copy(code"""
       |${CodeGenerator.javaType(dataType)} ${ev.value} = new $genericArrayData(new Object[0]);
       |boolean ${ev.isNull} = false;
     """.stripMargin)
   }
 
   def nonEmptyInputGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
-    val genericArrayData = classOf[GenericArrayData].getName
     val genericInternalRow = classOf[GenericInternalRow].getName
     val arrVals = ctx.freshName("arrVals")
     val arrCardinality = ctx.freshName("arrCardinality")
     val biggestCardinality = ctx.freshName("biggestCardinality")
 
-    val myobject = ctx.freshName("myobject")
+    val currentRow = ctx.freshName("currentRow")
     val j = ctx.freshName("j")
     val i = ctx.freshName("i")
     val args = ctx.freshName("args")
@@ -218,11 +217,11 @@ case class Zip(children: Seq[Expression]) extends Expression with ExpectsInputTy
     val getValueForType = arrayElementTypes.zipWithIndex.map { case (eleType, idx) =>
       val g = CodeGenerator.getValue(s"$arrVals[$idx]", eleType, i)
       s"""
-      |if ($i < $arrCardinality[$idx] && !$arrVals[$idx].isNullAt($i)) {
-      |  $myobject[$idx] = $g;
-      |} else {
-      |  $myobject[$idx] = null;
-      |}
+        |if ($i < $arrCardinality[$idx] && !$arrVals[$idx].isNullAt($i)) {
+        |  $currentRow[$idx] = $g;
+        |} else {
+        |  $currentRow[$idx] = null;
+        |}
       """.stripMargin
     }
 
@@ -231,7 +230,7 @@ case class Zip(children: Seq[Expression]) extends Expression with ExpectsInputTy
       funcName = "extractValue",
       arguments =
         ("int", i) ::
-        ("Object[]", myobject) ::
+        ("Object[]", currentRow) ::
         ("int[]", arrCardinality) ::
         ("ArrayData[]", arrVals) :: Nil)
 
@@ -249,9 +248,9 @@ case class Zip(children: Seq[Expression]) extends Expression with ExpectsInputTy
       |if (!${ev.isNull}) {
       |  Object[] $args = new Object[$biggestCardinality];
       |  for (int $i = 0; $i < $biggestCardinality; $i ++) {
-      |    Object[] $myobject = new Object[$numberOfArrays];
+      |    Object[] $currentRow = new Object[$numberOfArrays];
       |    $getValueForTypeSplitted
-      |    $args[$i] = new $genericInternalRow($myobject);
+      |    $args[$i] = new $genericInternalRow($currentRow);
       |  }
       |  ${ev.value} = new $genericArrayData($args);
       |}