diff --git a/sql/core/src/main/scala/org/apache/spark/sql/SchemaRDD.scala b/sql/core/src/main/scala/org/apache/spark/sql/SchemaRDD.scala
index 6aa4c7675674c..c7fac04034539 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/SchemaRDD.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/SchemaRDD.scala
@@ -314,11 +314,10 @@ class SchemaRDD(
   def analyze = sqlContext.analyzer(logicalPlan)
 
   def javaToPython: JavaRDD[Array[Byte]] = {
-    //val fieldNames: Seq[String] = logicalPlan.references.map(_.name)
+    val fieldNames: Seq[String] = this.queryExecution.analyzed.output.map(_.name)
     this.mapPartitions { iter =>
       val pickle = new Pickler
       iter.map { row =>
-        val fieldNames: Seq[String] = (1 to row.length).map(_.toString + "KEY") //TODO: Temporary
         val map: JMap[String, Any] = new java.util.HashMap
         val arr: java.util.ArrayList[Any] = new java.util.ArrayList
         row.zip(fieldNames).foreach { case (obj, name) =>