[SPARK-14098][SQL] Generate Java code that gets a float/double value in each column of CachedBatch when DataFrame.cache() is called

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

@@ -100,3 +100,28 @@ object BindReferences extends Logging {
     }.asInstanceOf[A] // Kind of a hack, but safe.  TODO: Tighten return type when possible.
   }
 }
+
+/**
+ * A column vector reference points to a specific column for ColumnVector.
+ * columnVar is a variable that keeps ColumnVector, and ordinal is row index in ColumnVector
+ */
+case class ColumnVectorReference(
+  columnVar: String, ordinal: String, dataType: DataType, nullable: Boolean)
+  extends LeafExpression {
+
+  override def toString: String = s"columnVector[$columnVar, $ordinal, ${dataType.simpleString}]"
+
+  override def eval(input: InternalRow): Any = null
+
+  override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+    val javaType = ctx.javaType(dataType)
+    val value = ctx.getValue(columnVar, dataType, ordinal)
+    if (nullable) {
+      ev.copy(code = s"""
+        boolean ${ev.isNull} = ${columnVar}.isNullAt($ordinal);
+        $javaType ${ev.value} = ${ev.isNull} ? ${ctx.defaultValue(dataType)} : ($value);""")
+    } else {
+      ev.copy(code = s"""$javaType ${ev.value} = $value;""", isNull = "false")
+    }
+  }
+}

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala

@@ -123,11 +123,12 @@ class CodegenContext {
    *
    * They will be kept as member variables in generated classes like `SpecificProjection`.
    */
-  val mutableStates: mutable.ArrayBuffer[(String, String, String)] =
-    mutable.ArrayBuffer.empty[(String, String, String)]
+  val mutableStates: mutable.ArrayBuffer[(String, String, String, String)] =
+    mutable.ArrayBuffer.empty[(String, String, String, String)]
 
-  def addMutableState(javaType: String, variableName: String, initCode: String): Unit = {
-    mutableStates += ((javaType, variableName, initCode))
+  def addMutableState(javaType: String, variableName: String, initCode: String,
+                      cleanupCode: String = ""): Unit = {
+    mutableStates += ((javaType, variableName, initCode, cleanupCode))
   }
 
   /**
@@ -149,7 +150,7 @@ class CodegenContext {
   def declareMutableStates(): String = {
     // It's possible that we add same mutable state twice, e.g. the `mergeExpressions` in
     // `TypedAggregateExpression`, we should call `distinct` here to remove the duplicated ones.
-    mutableStates.distinct.map { case (javaType, variableName, _) =>
+    mutableStates.distinct.map { case (javaType, variableName, _, _) =>
       s"private $javaType $variableName;"
     }.mkString("\n")
   }
@@ -160,6 +161,9 @@ class CodegenContext {
     mutableStates.distinct.map(_._3).mkString("\n")
   }
 
+  def cleanupMutableStates(): String = {
+    mutableStates.map(_._4).mkString("\n")
+  }
   /**
    * Holding all the functions those will be added into generated class.
    */
@@ -206,6 +210,10 @@ class CodegenContext {
   /** The variable name of the input row in generated code. */
   final var INPUT_ROW = "i"
 
+  var isRow = true
+  var enableColumnCodeGen = false
+  var iteratorInput = ""
+
   /**
    * The map from a variable name to it's next ID.
    */