Use TraversableOnce for regular Generators.

Author: hvanhovell

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

@@ -26,7 +26,6 @@ import com.fasterxml.jackson.core._
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.analysis.TypeCheckResult
 import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext, CodegenFallback, ExprCode}
-import org.apache.spark.sql.catalyst.util.GenericArrayData
 import org.apache.spark.sql.types.{DataType, StringType, StructField, StructType}
 import org.apache.spark.unsafe.types.UTF8String
 import org.apache.spark.util.Utils
@@ -328,9 +327,6 @@ case class GetJsonObject(json: Expression, path: Expression)
 // scalastyle:on line.size.limit
 case class JsonTuple(children: Seq[Expression]) extends Generator {
 
-  // a row is always returned
-  override def nullable: Boolean = false
-
   // if processing fails this shared value will be returned
   @transient private lazy val nullRow: Seq[InternalRow] =
     new GenericInternalRow(fieldExpressions.length) :: Nil
@@ -399,17 +395,16 @@ case class JsonTuple(children: Seq[Expression]) extends Generator {
   }
 
   override protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
-    val arrayDataClass = classOf[GenericArrayData].getName
+    val iteratorClass = classOf[Iterator[_]].getName
     val rowClass = classOf[GenericInternalRow].getName
-    def newArrayData(p: String): String = s"new $arrayDataClass(new Object[]{new $rowClass($p)})"
 
     // Add an empty row to default to.
     val fieldCount = fieldExpressions.length
     val nullRow = ctx.freshName("nullRow")
     ctx.addMutableState(
-      arrayDataClass,
+      rowClass,
       nullRow,
-      s"this.$nullRow = ${newArrayData(fieldCount.toString)};")
+      s"this.$nullRow = new $rowClass(${fieldCount.toString});")
 
     // Add the field names as a class field and add the foldable field names.
     val fieldNames = ctx.freshName("fieldNames")
@@ -432,19 +427,19 @@ case class JsonTuple(children: Seq[Expression]) extends Generator {
 
     // Create the generated code.
     val jsonSource = jsonExpr.genCode(ctx)
-    val result = ctx.freshName("result")
+    val raw = ctx.freshName("raw")
+    val row = ctx.freshName("row")
     val jsonTupleClass = classOf[JsonTuple].getName
     ev.copy(code = s"""
          |${jsonSource.code}
          |boolean ${ev.isNull} = false;
-         |ArrayData ${ev.value} = null;
-         |if (${jsonSource.isNull}) {
-         |  ${ev.value} = $nullRow;
-         |} else {
+         |InternalRow $row = $nullRow;
+         |if (!(${jsonSource.isNull})) {
          |  ${evalFieldNames.mkString("")}
-         |  Object[] $result = $jsonTupleClass.extractTuple(${jsonSource.value}, $fieldNames);
-         |  ${ev.value} = $result == null ? $nullRow : ${newArrayData(result)};
+         |  Object[] $raw = $jsonTupleClass.extractTuple(${jsonSource.value}, $fieldNames);
+         |  $row = $raw != null ? new $rowClass($raw) : $nullRow;
          |}
+         |$iteratorClass<InternalRow> ${ev.value} = $iteratorClass$$.MODULE$$.single($row);
     """.stripMargin)
   }
 }

sql/core/src/main/scala/org/apache/spark/sql/execution/GenerateExec.scala

@@ -113,46 +113,48 @@ case class GenerateExec(
   }
 
   override def doConsume(ctx: CodegenContext, input: Seq[ExprCode], row: ExprCode): String = {
+    ctx.currentVars = input
+    ctx.copyResult = true
+
+    // Add input rows to the values when we are joining
+    val values = if (join) {
+      input
+    } else {
+      Seq.empty
+    }
+
+    // Generate the driving expression.
+    val data = boundGenerator.genCode(ctx)
+
     boundGenerator match {
-      case e: Explode => codeGen(ctx, e.child, expand = false, input, row)
-      case g => codeGen(ctx, g, expand = true, input, row)
+      case e: Explode => codeGenExplode(ctx, e.child, values, data, row)
+      case g => codeGenTraversableOnce(ctx, g, values, data, row)
     }
   }
 
-  /** Generate code for Generate. */
-  private def codeGen(
+  /**
+   * Generate code for [[Explode]].
+   */
+  private def codeGenExplode(
       ctx: CodegenContext,
       e: Expression,
-      expand: Boolean,
       input: Seq[ExprCode],
+      data: ExprCode,
       row: ExprCode): String = {
-    ctx.currentVars = input
-    ctx.copyResult = true
-
-    // Generate the driving expression.
-    val data = e.genCode(ctx)
 
     // Generate looping variables.
-    val numOutput = metricTerm(ctx, "numOutputRows")
     val index = ctx.freshName("index")
-    val numElements = ctx.freshName("numElements")
 
     // Add a check if the generate outer flag is true.
     val checks = optionalCode(outer, data.isNull)
-    val (initArrayData, initValues, values) = e.dataType match {
-      case ArrayType(st: StructType, nullable) if expand =>
-        val rowCode = codeGenAccessor(ctx, data.value, "col", index, st, nullable, checks)
-        val extendedChecks = checks ++ optionalCode(nullable, rowCode.isNull)
-        val values = st.fields.toSeq.zipWithIndex.map { case (f, i) =>
-          codeGenAccessor(ctx, rowCode.value, f.name, s"$i", f.dataType, f.nullable, extendedChecks)
-        }
-        ("", rowCode.code, values)
 
+    // Generate code for either ArrayData or MapData
+    val (initMapData, values) = e.dataType match {
       case ArrayType(dataType, nullable) =>
-        ("", "", Seq(codeGenAccessor(ctx, data.value, "col", index, dataType, nullable, checks)))
+        ("", Seq(codeGenAccessor(ctx, data.value, "col", index, dataType, nullable, checks)))
 
       case MapType(keyType, valueType, valueContainsNull) =>
-        // Materialize the key and the value array before we enter the loop.
+        // Materialize the key and the value arrays before we enter the loop.
         val keyArray = ctx.freshName("keyArray")
         val valueArray = ctx.freshName("valueArray")
         val initArrayData =
@@ -163,34 +165,75 @@ case class GenerateExec(
         val values = Seq(
           codeGenAccessor(ctx, keyArray, "key", index, keyType, nullable = false, checks),
           codeGenAccessor(ctx, valueArray, "value", index, valueType, valueContainsNull, checks))
-        (initArrayData, "", values)
-    }
-
-    // Determine result vars.
-    val outputValues = if (join) {
-      input ++ values
-    } else {
-      values
+        (initArrayData, values)
     }
 
     // In case of outer we need to make sure the loop is executed at-least once when the array/map
     // contains no input. We do this by setting the looping index to -1 if there is no input,
     // evaluation of the array is prevented by a check in the accessor code.
+    val numElements = ctx.freshName("numElements")
     val init = if (outer) s"$numElements == 0 ? -1 : 0" else "0"
+    val numOutput = metricTerm(ctx, "numOutputRows")
     s"""
        |${data.code}
-       |$initArrayData
+       |$initMapData
        |int $numElements = ${data.isNull} ? 0 : ${data.value}.numElements();
        |for (int $index = $init; $index < $numElements; $index++) {
        |  $numOutput.add(1);
-       |  $initValues
-       |  ${consume(ctx, outputValues)}
+       |  ${consume(ctx, input ++ values)}
        |}
      """.stripMargin
   }
 
   /**
-   * Generate for accessor code for ArrayData and InternalRows.
+   * Generate code for a regular [[TraversableOnce]] returning [[Generator]].
+   */
+  private def codeGenTraversableOnce(
+      ctx: CodegenContext,
+      e: Expression,
+      input: Seq[ExprCode],
+      data: ExprCode,
+      row: ExprCode): String = {
+
+    // Generate looping variables.
+    val iterator = ctx.freshName("iterator")
+    val hasNext = ctx.freshName("hasNext")
+    val current = ctx.freshName("row")
+
+    // Add a check if the generate outer flag is true.
+    val checks = optionalCode(outer, hasNext)
+    val values = e.dataType match {
+      case ArrayType(st: StructType, nullable) =>
+        st.fields.toSeq.zipWithIndex.map { case (f, i) =>
+          codeGenAccessor(ctx, current, f.name, s"$i", f.dataType, f.nullable, checks)
+        }
+    }
+
+    // In case of outer we need to make sure the loop is executed at-least-once when the iterator
+    // contains no input. We do this by adding an 'outer' variable which guarantees execution of
+    // the first iteration even if there is no input. Evaluation of the iterator is prevented by a
+    // check in the accessor code.
+    val hasNextCode = s"$hasNext = $iterator.hasNext()"
+    val outerVal = ctx.freshName("outer")
+    def concatIfOuter(s1: String, s2: String): String = s1 + (if (outer) s2 else "")
+    val init = concatIfOuter(s"boolean $hasNextCode", s", $outerVal = true")
+    val check = concatIfOuter(hasNext, s"|| $outerVal")
+    val update = concatIfOuter(hasNextCode, s", $outerVal = false")
+    val next = if (outer) s"$hasNext ? $iterator.next() : null" else s"$iterator.next()"
+    val numOutput = metricTerm(ctx, "numOutputRows")
+    s"""
+       |${data.code}
+       |scala.collection.Iterator<InternalRow> $iterator = ${data.value}.toIterator();
+       |for ($init; $check; $update) {
+       |  $numOutput.add(1);
+       |  InternalRow $current = (InternalRow) $next;
+       |  ${consume(ctx, input ++ values)}
+       |}
+     """.stripMargin
+  }
+
+  /**
+   * Generate accessor code for ArrayData and InternalRows.
    */
   private def codeGenAccessor(
       ctx: CodegenContext,
@@ -210,7 +253,7 @@ case class GenerateExec(
         s"""
            |boolean $isNull = ${checks.mkString(" || ")};
            |$javaType $value = $isNull ? ${ctx.defaultValue(dt)} : $getter;
-           """.stripMargin
+         """.stripMargin
       ExprCode(code, isNull, value)
     } else {
       ExprCode(s"$javaType $value = $getter;", "false", value)